From 04a30fb375a6c1312365fbbbdd0a7d0669212e92 Mon Sep 17 00:00:00 2001 From: Terry Jia Date: Thu, 9 Jul 2026 10:42:20 -0400 Subject: [PATCH 01/37] fix: Load3D failing path validation from double path resolution (#14852) --- comfy_extras/nodes_load_3d.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/comfy_extras/nodes_load_3d.py b/comfy_extras/nodes_load_3d.py index 6e3e88471..6ef9a1ca3 100644 --- a/comfy_extras/nodes_load_3d.py +++ b/comfy_extras/nodes_load_3d.py @@ -61,14 +61,10 @@ class Load3D(IO.ComfyNode): @classmethod def execute(cls, model_file, image, **kwargs) -> IO.NodeOutput: - image_path = folder_paths.get_annotated_filepath(image['image']) - mask_path = folder_paths.get_annotated_filepath(image['mask']) - normal_path = folder_paths.get_annotated_filepath(image['normal']) - load_image_node = nodes.LoadImage() - output_image, ignore_mask = load_image_node.load_image(image=image_path) - ignore_image, output_mask = load_image_node.load_image(image=mask_path) - normal_image, ignore_mask2 = load_image_node.load_image(image=normal_path) + output_image, ignore_mask = load_image_node.load_image(image=image['image']) + ignore_image, output_mask = load_image_node.load_image(image=image['mask']) + normal_image, ignore_mask2 = load_image_node.load_image(image=image['normal']) video = None From 412aaab0e27d244f3fe47b2e593fa40cc9bcb4c6 Mon Sep 17 00:00:00 2001 From: Simon Pinfold Date: Fri, 10 Jul 2026 07:59:30 +1200 Subject: [PATCH 02/37] feat(api): expose registered extension filters on /experiment/models (#14797) Each folder in the listing now carries its registered extension allowlist verbatim; an empty array means the folder accepts any extension (match-all), mirroring filter_files_extensions semantics. Gives consumers the filtering rule itself rather than just its output: /models/{folder} lists files by the per-folder rule but the rule is not exposed anywhere, and /experiment/models/{folder} filters everything by the global supported_pt_extensions regardless of registration. Presentation-level filtering of match-all folders (e.g. hiding README/config noise that repository-downloading custom nodes leave in model directories) is deliberately left to the consumer. Co-authored-by: guill --- app/model_manager.py | 6 +++++- openapi.yaml | 8 ++++++++ tests-unit/app_test/model_manager_test.py | 22 ++++++++++++++++++++++ 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/app/model_manager.py b/app/model_manager.py index b0329ce17..5928781ca 100644 --- a/app/model_manager.py +++ b/app/model_manager.py @@ -35,7 +35,11 @@ class ModelFileManager: for folder in model_types: if folder in folder_black_list: continue - output_folders.append({"name": folder, "folders": folder_paths.get_folder_paths(folder)}) + output_folders.append({ + "name": folder, + "folders": folder_paths.get_folder_paths(folder), + "extensions": sorted(folder_paths.folder_names_and_paths[folder][1]), + }) return web.json_response(output_folders) # NOTE: This is an experiment to replace `/models/{folder}` diff --git a/openapi.yaml b/openapi.yaml index 0cf177815..c09b1eeac 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -775,6 +775,14 @@ components: ModelFolder: description: Represents a folder containing models properties: + extensions: + description: The folder's registered file-extension allowlist. An empty array means the folder accepts any extension (match-all). + example: + - .ckpt + - .safetensors + items: + type: string + type: array folders: description: List of paths where models of this type are stored example: diff --git a/tests-unit/app_test/model_manager_test.py b/tests-unit/app_test/model_manager_test.py index ae59206f6..d7cc20fcd 100644 --- a/tests-unit/app_test/model_manager_test.py +++ b/tests-unit/app_test/model_manager_test.py @@ -24,6 +24,28 @@ def app(model_manager): app.add_routes(routes) return app +async def test_get_model_folders_includes_registered_extensions(aiohttp_client, app, tmp_path): + """Folders expose their registered extension set verbatim; an empty list + means match-all (filter_files_extensions semantics).""" + with patch('folder_paths.folder_names_and_paths', { + 'test_checkpoints': ([str(tmp_path)], {'.safetensors', '.ckpt'}), + 'test_configs': ([str(tmp_path)], ['.yaml']), + 'test_match_all': ([str(tmp_path)], set()), + 'configs': ([str(tmp_path)], ['.yaml']), + }): + client = await aiohttp_client(app) + response = await client.get('/experiment/models') + + assert response.status == 200 + folders = {f['name']: f for f in await response.json()} + + assert 'configs' not in folders # blocklisted + assert folders['test_checkpoints']['folders'] == [str(tmp_path)] + assert folders['test_checkpoints']['extensions'] == ['.ckpt', '.safetensors'] + assert folders['test_configs']['extensions'] == ['.yaml'] + # Match-all registrations are exposed honestly, not substituted. + assert folders['test_match_all']['extensions'] == [] + async def test_get_model_preview_safetensors(aiohttp_client, app, tmp_path): img = Image.new('RGB', (100, 100), 'white') img_byte_arr = BytesIO() From 1ea724339ca6bea1808e6600fefdd08bf42822e6 Mon Sep 17 00:00:00 2001 From: Alexis Rolland Date: Fri, 10 Jul 2026 05:57:44 +0800 Subject: [PATCH 03/37] Update cla.yml (#14851) --- .github/workflows/cla.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/cla.yml b/.github/workflows/cla.yml index b75397e50..bc0f779cf 100644 --- a/.github/workflows/cla.yml +++ b/.github/workflows/cla.yml @@ -32,9 +32,11 @@ jobs: PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }} PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }} BASE_ALLOWLIST: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot] + # For each commit emit the GitHub login when the author/committer email resolves to a GitHub account + # otherwise fall back to the raw git name. run: | others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \ - --jq '.[] | (.author.login // empty), (.committer.login // empty)' \ + --jq '.[] | (.author.login // .commit.author.name // empty), (.committer.login // .commit.committer.name // empty)' \ | sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -) if [ -n "$others" ]; then echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT" @@ -43,7 +45,7 @@ jobs: fi - name: CLA Assistant - # Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase. + # Run on PR events, on "recheck" comment, or when someone posts the signing phrase. # IMPORTANT: this phrase must match `custom-pr-sign-comment` below. if: > github.event_name == 'pull_request_target' || From 73e84d5ec8b943dcb42535229eb94ee7ab3abea1 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:57:09 -0700 Subject: [PATCH 04/37] Support convrot int4 models. (#14859) linear_dtype in comfy_quant metadata can be used to set if the int4 op does the matrix multiplication in int8 or int4, the default is int4 on GPUs that support it with fallback to int8 for GPUs that don't. --- comfy/ops.py | 26 +++++++++ comfy/quant_ops.py | 14 +++++ requirements.txt | 2 +- .../comfy_quant/test_mixed_precision.py | 54 ++++++++++++++++++- 4 files changed, 94 insertions(+), 2 deletions(-) diff --git a/comfy/ops.py b/comfy/ops.py index 35a1ee31e..0c6fe4cb4 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1104,6 +1104,21 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat scales["convrot_groupsize"] = int( layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256)) ) + elif module.quant_format == "convrot_w4a4": + scale = pop_scale("weight_scale") + if scale is None: + raise ValueError(f"Missing ConvRot W4A4 weight scale for layer {layer_name}") + params_conf = layer_conf.get("params", {}) + if not isinstance(params_conf, dict): + params_conf = {} + scales = { + "scale": scale, + "convrot_groupsize": int( + layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256)) + ), + "quant_group_size": 64, + "linear_dtype": layer_conf.get("linear_dtype", params_conf.get("linear_dtype", "int4")), + } else: raise ValueError(f"Unsupported quantization format: {module.quant_format}") @@ -1150,6 +1165,11 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False): quant_conf["convrot"] = True quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256) + elif module.quant_format == "convrot_w4a4": + quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256) + linear_dtype = getattr(params, "linear_dtype", "int4") + if linear_dtype != "int4": + quant_conf["linear_dtype"] = linear_dtype if extra_quant_conf: quant_conf.update(extra_quant_conf) sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8) @@ -1430,6 +1450,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec } if hasattr(params, "block_scale"): # NVFP4 kwargs["block_scale"] = params.block_scale[i] + if hasattr(params, "quant_group_size"): + kwargs["quant_group_size"] = params.quant_group_size + if hasattr(params, "convrot_groupsize"): + kwargs["convrot_groupsize"] = params.convrot_groupsize + if hasattr(params, "linear_dtype"): + kwargs["linear_dtype"] = params.linear_dtype return QuantizedTensor(weight._qdata[i], weight._layout_cls, type(params)(**kwargs)) def state_dict(self, *args, destination=None, prefix="", **kwargs): diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py index 44f25a97e..53a0cb603 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -10,6 +10,7 @@ try: QuantizedLayout, TensorCoreFP8Layout as _CKFp8Layout, TensorCoreNVFP4Layout as _CKNvfp4Layout, + TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout, TensorWiseINT8Layout as _CKTensorWiseINT8Layout, register_layout_op, register_layout_class, @@ -51,6 +52,9 @@ except ImportError as e: class _CKTensorWiseINT8Layout: pass + class _CKTensorCoreConvRotW4A4Layout: + pass + def register_layout_class(name, cls): pass @@ -179,6 +183,7 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase): # Backward compatibility alias - default to E4M3 TensorCoreFP8Layout = TensorCoreFP8E4M3Layout TensorWiseINT8Layout = _CKTensorWiseINT8Layout +TensorCoreConvRotW4A4Layout = _CKTensorCoreConvRotW4A4Layout # ============================================================================== @@ -190,6 +195,7 @@ register_layout_class("TensorCoreFP8E4M3Layout", TensorCoreFP8E4M3Layout) register_layout_class("TensorCoreFP8E5M2Layout", TensorCoreFP8E5M2Layout) register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout) register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout) +register_layout_class("TensorCoreConvRotW4A4Layout", _CKTensorCoreConvRotW4A4Layout) if _CK_MXFP8_AVAILABLE: register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout) @@ -227,6 +233,13 @@ QUANT_ALGOS["int8_tensorwise"] = { "quantize_input": False, } +QUANT_ALGOS["convrot_w4a4"] = { + "storage_t": torch.int8, + "parameters": {"weight_scale"}, + "comfy_tensor_layout": "TensorCoreConvRotW4A4Layout", + "quantize_input": False, +} + # ============================================================================== # Re-exports for backward compatibility @@ -239,6 +252,7 @@ __all__ = [ "TensorCoreFP8E4M3Layout", "TensorCoreFP8E5M2Layout", "TensorCoreNVFP4Layout", + "TensorCoreConvRotW4A4Layout", "TensorWiseINT8Layout", "QUANT_ALGOS", "register_layout_op", diff --git a/requirements.txt b/requirements.txt index e72f3045b..a8ea0eace 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,7 +22,7 @@ alembic SQLAlchemy>=2.0.0 filelock av>=16.0.0 -comfy-kitchen==0.2.16 +comfy-kitchen==0.2.17 comfy-aimdo==0.4.10 requests simpleeval>=1.0.0 diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py index 43b4b7ce9..7bbc96616 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/tests-unit/comfy_quant/test_mixed_precision.py @@ -15,7 +15,7 @@ if not has_gpu(): args.cpu = True from comfy import ops -from comfy.quant_ops import QuantizedTensor +from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor import comfy.utils @@ -283,7 +283,59 @@ class TestMixedPrecisionOps(unittest.TestCase): saved = model.state_dict() saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes()) self.assertTrue(saved_conf["convrot"]) + + def test_convrot_w4a4_loads_into_params(self): + """ConvRot W4A4 checkpoints must load as the dedicated kitchen layout.""" + if "convrot_w4a4" not in QUANT_ALGOS: + self.skipTest("comfy_kitchen does not provide ConvRot W4A4") + + torch.manual_seed(456) + layer_quant_config = { + "layer": { + "format": "convrot_w4a4", + "convrot_groupsize": 256, + "linear_dtype": "int8", + } + } + weight = torch.randn(16, 256, dtype=torch.bfloat16) + bias = torch.randn(16, dtype=torch.bfloat16) + q_weight = QuantizedTensor.from_float( + weight, + "TensorCoreConvRotW4A4Layout", + convrot_groupsize=256, + quant_group_size=64, + ) + state_dict = { + "layer.weight": q_weight._qdata, + "layer.bias": bias, + "layer.weight_scale": q_weight._params.scale, + } + + state_dict, _ = comfy.utils.convert_old_quants( + state_dict, + metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}, + ) + model = torch.nn.Module() + model.layer = ops.mixed_precision_ops({}).Linear(256, 16, device="cpu", dtype=torch.bfloat16) + model.load_state_dict(state_dict, strict=False) + + self.assertIsInstance(model.layer.weight, QuantizedTensor) + self.assertEqual(model.layer.weight._layout_cls, "TensorCoreConvRotW4A4Layout") + self.assertEqual(model.layer.weight._params.convrot_groupsize, 256) + self.assertEqual(model.layer.weight._params.quant_group_size, 64) + self.assertEqual(model.layer.weight._params.linear_dtype, "int8") + + input_tensor = torch.randn(4, 256, dtype=torch.bfloat16) + loaded_out = model.layer(input_tensor) + ref_out = torch.nn.functional.linear(input_tensor, q_weight, bias) + self.assertTrue(torch.equal(loaded_out, ref_out)) + + saved = model.state_dict() + saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes()) + self.assertEqual(saved_conf["format"], "convrot_w4a4") self.assertEqual(saved_conf["convrot_groupsize"], 256) + self.assertEqual(saved_conf["linear_dtype"], "int8") + self.assertNotIn("quant_group_size", saved_conf) if __name__ == "__main__": unittest.main() From b7a648ca2011489ba40eaacf01a5d6f4e9fab539 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:39:01 -0700 Subject: [PATCH 05/37] Try to fix the model reloading issue some people have. (#14822) --- comfy/model_management.py | 11 +++++++++++ comfy_execution/caching.py | 25 +++++++++++++++++-------- execution.py | 17 +++++++++++------ 3 files changed, 39 insertions(+), 14 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index b15d08ba1..222005b6f 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -616,6 +616,8 @@ 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 +WINDOWS_PIN_EVICTION_SWAP_PERCENT = 5.0 +WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE = 512 * 1024 ** 2 def module_size(module): module_mem = 0 @@ -642,6 +644,15 @@ def free_pins(size, evict_active=False): size -= freed return freed_total +def should_free_pins_for_ram_pressure(shortfall): + if shortfall <= 0: + return False + if not WINDOWS: + return True + if psutil.virtual_memory().available < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE: + return True + return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT + def ensure_pin_budget(size, evict_active=False): if args.high_ram: return True diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index ad75a0e50..6bd99b68f 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -503,6 +503,8 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05 RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3 +RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2 + def all_outputs_dynamic(outputs): if outputs is None: @@ -517,7 +519,6 @@ def all_outputs_dynamic(outputs): return True - class RAMPressureCache(LRUCache): def __init__(self, key_class, enable_providers=False): @@ -539,9 +540,9 @@ class RAMPressureCache(LRUCache): self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time() super().set_local(node_id, value) - def ram_release(self, target, free_active=False): + def ram_release(self, target, free_active=False, min_entry_size=0): if psutil.virtual_memory().available >= target: - return + return 0 clean_list = [] @@ -555,8 +556,9 @@ class RAMPressureCache(LRUCache): oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key]) ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE + oom_ram_usage = ram_usage def scan_list_for_ram_usage(outputs): - nonlocal ram_usage + nonlocal ram_usage, oom_ram_usage if outputs is None: return for output in outputs: @@ -564,19 +566,26 @@ class RAMPressureCache(LRUCache): scan_list_for_ram_usage(output) elif isinstance(output, torch.Tensor) and output.device.type == 'cpu': ram_usage += output.numel() * output.element_size() + oom_ram_usage += output.numel() * output.element_size() elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation: #old ModelPatchers are the first to go - ram_usage = 1e30 + oom_ram_usage = 1e30 scan_list_for_ram_usage(cache_entry.outputs) - oom_score *= ram_usage + if ram_usage < min_entry_size: + continue + + oom_score *= oom_ram_usage #In the case where we have no information on the node ram usage at all, #break OOM score ties on the last touch timestamp (pure LRU) - bisect.insort(clean_list, (oom_score, self.timestamps[key], key)) + bisect.insort(clean_list, (oom_score, self.timestamps[key], key, ram_usage)) + freed = 0 while psutil.virtual_memory().available < target and clean_list: - _, _, key = clean_list.pop() + _, _, key, ram_usage = clean_list.pop() del self.cache[key] self.used_generation.pop(key, None) self.timestamps.pop(key, None) self.children.pop(key, None) + freed += ram_usage + return freed diff --git a/execution.py b/execution.py index c45317593..19b8cdd68 100644 --- a/execution.py +++ b/execution.py @@ -29,6 +29,7 @@ from comfy_execution.caching import ( HierarchicalCache, LRUCache, RAMPressureCache, + RAM_CACHE_LARGE_INTERMEDIATE, ) from comfy_execution.graph import ( DynamicPrompt, @@ -794,12 +795,16 @@ class PromptExecutor: 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) + if ram_shortfall > 0: + freed = ram_release_callback(ram_headroom, free_active=True, min_entry_size=RAM_CACHE_LARGE_INTERMEDIATE) + ram_shortfall -= freed + if comfy.model_management.should_free_pins_for_ram_pressure(ram_shortfall): + 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) else: # Only execute when the while-loop ends without break # Send cached UI for intermediate output nodes that weren't executed From 62e025a4f34d16eeedfc1c93e50a48a69098df7f Mon Sep 17 00:00:00 2001 From: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:30:26 +0800 Subject: [PATCH 06/37] Fix FP8 activation quantization for >2D activations in mixed_precision_ops (#14643) mixed_precision_ops.Linear.forward only quantized activations that were 2D, or 3D (reshaped to 2D). Inputs with rank >= 4 (e.g. Anima's MLP activations, which are not reshaped to 3D the way the attention path is) fell through the `input_reshaped.ndim == 2` guard and reached scaled_mm as bf16, silently dispatching a bf16 kernel instead of FP8. Since MLP is roughly half the compute, the FP8 speedup was far below expectation. Generalize the existing 3D->2D reshape to any rank >= 3 (flatten the leading dims, keep the contraction dim) and reshape the output back to the original leading dims. 2D and 3D inputs are handled exactly as before; only rank >= 4 inputs change (now quantized instead of skipped). This matches the rank-agnostic handling already used by the training path (flatten(0, -2) / unflatten). Fixes #14595. --- comfy/ops.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/comfy/ops.py b/comfy/ops.py index 0c6fe4cb4..13c2604fb 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1257,7 +1257,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec run_every_op() input_shape = input.shape - reshaped_3d = False + reshaped_nd = False #If cast needs to apply lora, it should be done in the compute dtype compute_dtype = input.dtype @@ -1294,12 +1294,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec # Inference path (unchanged) if _use_quantized and quantize_input: - # Reshape 3D tensors to 2D for quantization (needed for NVFP4 and others) - input_reshaped = input.reshape(-1, input_shape[2]) if input.ndim == 3 else input + # Reshape >=3D tensors to 2D for quantization (needed for NVFP4 and others) + input_reshaped = input.reshape(-1, input_shape[-1]) if input.ndim >= 3 else input # Fall back to non-quantized for non-2D tensors if input_reshaped.ndim == 2: - reshaped_3d = input.ndim == 3 + reshaped_nd = input.ndim >= 3 # dtype is now implicit in the layout class scale = getattr(self, 'input_scale', None) if scale is not None: @@ -1314,9 +1314,9 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec weight_only_quant=weight_only_quant, ) - # Reshape output back to 3D if input was 3D - if reshaped_3d: - output = output.reshape((input_shape[0], input_shape[1], self.weight.shape[0])) + # Reshape output back to original rank if input was >2D + if reshaped_nd: + output = output.reshape((*input_shape[:-1], self.weight.shape[0])) return output From 099522f85bcd8586eac4133c02f31c70dafe85d2 Mon Sep 17 00:00:00 2001 From: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> Date: Fri, 10 Jul 2026 11:11:52 +0800 Subject: [PATCH 07/37] Enable comfy-kitchen Triton backend by default on ROCm/AMD (#14862) On AMD/ROCm the CUDA backend is unavailable, so Triton is the only accelerated comfy-kitchen backend. It was disabled by default (opt-in --enable-triton-backend), leaving AMD on the slow eager path. Enable it by default when torch.version.hip is set AND Triton is >= 3.7 -- older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path, so on Triton < 3.7 it stays disabled with a log line. NVIDIA behavior is unchanged; the explicit --enable-triton-backend flag still works as an override. Fixes #14861 --- comfy/quant_ops.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py index 53a0cb603..91b3e4fe9 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -25,10 +25,18 @@ try: ck.registry.disable("cuda") logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.") - if args.enable_triton_backend: + # On ROCm/AMD the CUDA backend is unavailable, so Triton is the only accelerated + # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7: + # older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path. + if args.enable_triton_backend or torch.version.hip is not None: try: import triton - logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__) + triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2]) + if args.enable_triton_backend or triton_version >= (3, 7): + logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__) + else: + logging.info("Triton %s is too old for the ROCm INT8 path (needs >= 3.7); comfy-kitchen triton backend disabled.", triton.__version__) + ck.registry.disable("triton") except ImportError as e: logging.error(f"Failed to import triton, Error: {e}, the comfy-kitchen triton backend will not be available.") ck.registry.disable("triton") From e2a6e30d892402ffcf01d6280c8e2744a4448b9d Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Thu, 9 Jul 2026 20:17:06 -0700 Subject: [PATCH 08/37] Fix black image on turing when using int4 models. (#14864) --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index a8ea0eace..790ef4940 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,7 +22,7 @@ alembic SQLAlchemy>=2.0.0 filelock av>=16.0.0 -comfy-kitchen==0.2.17 +comfy-kitchen==0.2.18 comfy-aimdo==0.4.10 requests simpleeval>=1.0.0 From 8e2e54e2b8fd2623e5bc255b7c90ccf09bcb04bb Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 10 Jul 2026 02:07:42 -0500 Subject: [PATCH 09/37] Add SeedVR2 support (CORE-6) (#14424) --- comfy/latent_formats.py | 4 + comfy/ldm/modules/diffusionmodules/model.py | 6 +- comfy/ldm/seedvr/attention.py | 51 + comfy/ldm/seedvr/color_fix.py | 301 +++ comfy/ldm/seedvr/constants.py | 48 + comfy/ldm/seedvr/model.py | 1361 ++++++++++++++ comfy/ldm/seedvr/vae.py | 1612 +++++++++++++++++ comfy/model_base.py | 12 + comfy/model_detection.py | 45 +- comfy/sd.py | 102 +- comfy/supported_models.py | 35 + comfy/supported_models_base.py | 6 +- comfy/text_encoders/gemma4.py | 2 +- comfy_extras/nodes_seedvr.py | 614 +++++++ nodes.py | 1 + .../test_seedvr2_conditioning.py | 186 ++ .../comfy_extras_test/test_seedvr2_nodes.py | 55 + .../test_seedvr2_post_processing.py | 51 + .../test_seedvr2_temporal_chunk.py | 77 + tests-unit/comfy_test/model_detection_test.py | 83 +- .../comfy_test/seedvr_vae_forward_test.py | 74 + tests-unit/comfy_test/test_seedvr2_dtype.py | 50 + .../comfy_test/test_seedvr2_internals.py | 169 ++ tests-unit/comfy_test/test_seedvr2_model.py | 320 ++++ .../comfy_test/test_seedvr2_vae_decode.py | 94 + .../comfy_test/test_seedvr2_vae_tiled.py | 382 ++++ 26 files changed, 5712 insertions(+), 29 deletions(-) create mode 100644 comfy/ldm/seedvr/attention.py create mode 100644 comfy/ldm/seedvr/color_fix.py create mode 100644 comfy/ldm/seedvr/constants.py create mode 100644 comfy/ldm/seedvr/model.py create mode 100644 comfy/ldm/seedvr/vae.py create mode 100644 comfy_extras/nodes_seedvr.py create mode 100644 tests-unit/comfy_extras_test/test_seedvr2_conditioning.py create mode 100644 tests-unit/comfy_extras_test/test_seedvr2_nodes.py create mode 100644 tests-unit/comfy_extras_test/test_seedvr2_post_processing.py create mode 100644 tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py create mode 100644 tests-unit/comfy_test/seedvr_vae_forward_test.py create mode 100644 tests-unit/comfy_test/test_seedvr2_dtype.py create mode 100644 tests-unit/comfy_test/test_seedvr2_internals.py create mode 100644 tests-unit/comfy_test/test_seedvr2_model.py create mode 100644 tests-unit/comfy_test/test_seedvr2_vae_decode.py create mode 100644 tests-unit/comfy_test/test_seedvr2_vae_tiled.py diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index bbdfd4bc2..8a16cfe55 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -779,6 +779,10 @@ class ACEAudio(LatentFormat): latent_channels = 8 latent_dimensions = 2 +class SeedVR2(LatentFormat): + latent_channels = 16 + latent_dimensions = 3 + class ACEAudio15(LatentFormat): latent_channels = 64 latent_dimensions = 1 diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py index fcbaa074f..e752d0ecb 100644 --- a/comfy/ldm/modules/diffusionmodules/model.py +++ b/comfy/ldm/modules/diffusionmodules/model.py @@ -22,7 +22,7 @@ def torch_cat_if_needed(xl, dim): else: return None -def get_timestep_embedding(timesteps, embedding_dim): +def get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos=False, downscale_freq_shift=1): """ This matches the implementation in Denoising Diffusion Probabilistic Models: From Fairseq. @@ -33,11 +33,13 @@ def get_timestep_embedding(timesteps, embedding_dim): assert len(timesteps.shape) == 1 half_dim = embedding_dim // 2 - emb = math.log(10000) / (half_dim - 1) + emb = math.log(10000) / (half_dim - downscale_freq_shift) emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) emb = emb.to(device=timesteps.device) emb = timesteps.float()[:, None] * emb[None, :] emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) if embedding_dim % 2 == 1: # zero pad emb = torch.nn.functional.pad(emb, (0,1,0,0)) return emb diff --git a/comfy/ldm/seedvr/attention.py b/comfy/ldm/seedvr/attention.py new file mode 100644 index 000000000..11b4c1e4a --- /dev/null +++ b/comfy/ldm/seedvr/attention.py @@ -0,0 +1,51 @@ +import torch + +from comfy.ldm.modules import attention as _attention + + +def _var_attention_qkv(q, k, v, heads, skip_reshape): + if skip_reshape: + return q, k, v, q.shape[-1] + total_tokens, embed_dim = q.shape + head_dim = embed_dim // heads + return ( + q.view(total_tokens, heads, head_dim), + k.view(k.shape[0], heads, head_dim), + v.view(v.shape[0], heads, head_dim), + head_dim, + ) + + +def _var_attention_output(out, heads, head_dim, skip_output_reshape): + if skip_output_reshape: + return out + return out.reshape(-1, heads * head_dim) + + +def var_attention_optimized_split(q, k, v, heads, cu_seqlens_q, cu_seqlens_k, *args, skip_reshape=False, skip_output_reshape=False, **kwargs): + q, k, v, head_dim = _var_attention_qkv(q, k, v, heads, skip_reshape) + + q_split_indices = cu_seqlens_q[1:-1] + k_split_indices = cu_seqlens_k[1:-1] + if k.shape[0] != v.shape[0]: + raise ValueError("cu_seqlens_k does not match v token count") + + q_splits = torch.tensor_split(q, q_split_indices, dim=0) + k_splits = torch.tensor_split(k, k_split_indices, dim=0) + v_splits = torch.tensor_split(v, k_split_indices, dim=0) + if len(q_splits) != len(k_splits) or len(q_splits) != len(v_splits): + raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same sequence count") + + out = [] + for q_i, k_i, v_i in zip(q_splits, k_splits, v_splits): + q_i = q_i.permute(1, 0, 2).unsqueeze(0) + k_i = k_i.permute(1, 0, 2).unsqueeze(0) + v_i = v_i.permute(1, 0, 2).unsqueeze(0) + out_i = _attention.optimized_attention(q_i, k_i, v_i, heads, skip_reshape=True, skip_output_reshape=True) + out.append(out_i.squeeze(0).permute(1, 0, 2)) + + out = torch.cat(out, dim=0) + return _var_attention_output(out, heads, head_dim, skip_output_reshape) + + +optimized_var_attention = var_attention_optimized_split diff --git a/comfy/ldm/seedvr/color_fix.py b/comfy/ldm/seedvr/color_fix.py new file mode 100644 index 000000000..a43cb5270 --- /dev/null +++ b/comfy/ldm/seedvr/color_fix.py @@ -0,0 +1,301 @@ +import torch +import torch.nn.functional as F +from torch import Tensor + +from comfy.ldm.seedvr.constants import ( + CIELAB_DELTA, + CIELAB_KAPPA, + D65_WHITE_X, + D65_WHITE_Z, + WAVELET_DECOMP_LEVELS, +) + + +def wavelet_blur(image: Tensor, radius): + max_safe_radius = max(1, min(image.shape[-2:]) // 8) + if radius > max_safe_radius: + radius = max_safe_radius + + num_channels = image.shape[1] + + kernel_vals = [ + [0.0625, 0.125, 0.0625], + [0.125, 0.25, 0.125], + [0.0625, 0.125, 0.0625], + ] + kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device) + kernel = kernel[None, None].repeat(num_channels, 1, 1, 1) + + image = F.pad(image, (radius, radius, radius, radius), mode='replicate') + output = F.conv2d(image, kernel, groups=num_channels, dilation=radius) + + return output + +def wavelet_decomposition(image: Tensor, levels: int = WAVELET_DECOMP_LEVELS): + high_freq = torch.zeros_like(image) + + for i in range(levels): + radius = 2 ** i + low_freq = wavelet_blur(image, radius) + high_freq.add_(image).sub_(low_freq) + image = low_freq + + return high_freq, low_freq + +def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor) -> Tensor: + + if content_feat.shape != style_feat.shape: + if len(content_feat.shape) >= 3: + style_feat = F.interpolate( + style_feat, + size=content_feat.shape[-2:], + mode='bilinear', + align_corners=False + ) + + content_high_freq, content_low_freq = wavelet_decomposition(content_feat) + del content_low_freq + + style_high_freq, style_low_freq = wavelet_decomposition(style_feat) + del style_high_freq + + if content_high_freq.shape != style_low_freq.shape: + style_low_freq = F.interpolate( + style_low_freq, + size=content_high_freq.shape[-2:], + mode='bilinear', + align_corners=False + ) + + content_high_freq.add_(style_low_freq) + + return content_high_freq.clamp_(-1.0, 1.0) + +def _histogram_matching_channel(source: Tensor, reference: Tensor) -> Tensor: + original_shape = source.shape + + source_flat = source.flatten() + reference_flat = reference.flatten() + + source_sorted, source_indices = torch.sort(source_flat) + reference_sorted, _ = torch.sort(reference_flat) + del reference_flat + + n_source = len(source_sorted) + n_reference = len(reference_sorted) + + if n_source == n_reference: + matched_sorted = reference_sorted + else: + source_quantiles = torch.linspace(0, 1, n_source, device=source.device) + ref_indices = (source_quantiles * (n_reference - 1)).long() + ref_indices.clamp_(0, n_reference - 1) + matched_sorted = reference_sorted[ref_indices] + del source_quantiles, ref_indices, reference_sorted + + del source_sorted, source_flat + + inverse_indices = torch.argsort(source_indices) + del source_indices + matched_flat = matched_sorted[inverse_indices] + del matched_sorted, inverse_indices + + return matched_flat.reshape(original_shape) + +def _lab_to_rgb_batch(lab: Tensor, matrix_inv: Tensor, epsilon: float, kappa: float) -> Tensor: + L, a, b = lab[:, 0], lab[:, 1], lab[:, 2] + + fy = (L + 16.0) / 116.0 + fx = a.div(500.0).add_(fy) + fz = fy - b / 200.0 + del L, a, b + + x = torch.where( + fx > epsilon, + torch.pow(fx, 3.0), + fx.mul(116.0).sub_(16.0).div_(kappa) + ) + y = torch.where( + fy > epsilon, + torch.pow(fy, 3.0), + fy.mul(116.0).sub_(16.0).div_(kappa) + ) + z = torch.where( + fz > epsilon, + torch.pow(fz, 3.0), + fz.mul(116.0).sub_(16.0).div_(kappa) + ) + del fx, fy, fz + + x.mul_(D65_WHITE_X) + z.mul_(D65_WHITE_Z) + + xyz = torch.stack([x, y, z], dim=1) + del x, y, z + + B, _, H, W = xyz.shape + xyz_flat = xyz.permute(0, 2, 3, 1).reshape(-1, 3) + del xyz + + xyz_flat = xyz_flat.to(dtype=matrix_inv.dtype) + rgb_linear_flat = torch.matmul(xyz_flat, matrix_inv.T) + del xyz_flat + + rgb_linear = rgb_linear_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2) + del rgb_linear_flat + + mask = rgb_linear > 0.0031308 + rgb = torch.where( + mask, + torch.pow(torch.clamp(rgb_linear, min=0.0), 1.0 / 2.4).mul_(1.055).sub_(0.055), + rgb_linear * 12.92 + ) + del mask, rgb_linear + + return torch.clamp(rgb, 0.0, 1.0) + +def _rgb_to_lab_batch(rgb: Tensor, matrix: Tensor, epsilon: float, kappa: float) -> Tensor: + mask = rgb > 0.04045 + rgb_linear = torch.where( + mask, + torch.pow((rgb + 0.055) / 1.055, 2.4), + rgb / 12.92 + ) + del mask + + B, _, H, W = rgb_linear.shape + rgb_flat = rgb_linear.permute(0, 2, 3, 1).reshape(-1, 3) + del rgb_linear + + rgb_flat = rgb_flat.to(dtype=matrix.dtype) + xyz_flat = torch.matmul(rgb_flat, matrix.T) + del rgb_flat + + xyz = xyz_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2) + del xyz_flat + + xyz[:, 0].div_(D65_WHITE_X) + xyz[:, 2].div_(D65_WHITE_Z) + + epsilon_cubed = epsilon ** 3 + mask = xyz > epsilon_cubed + f_xyz = torch.where( + mask, + torch.pow(xyz, 1.0 / 3.0), + xyz.mul(kappa).add_(16.0).div_(116.0) + ) + del xyz, mask + + L = f_xyz[:, 1].mul(116.0).sub_(16.0) + a = (f_xyz[:, 0] - f_xyz[:, 1]).mul_(500.0) + b = (f_xyz[:, 1] - f_xyz[:, 2]).mul_(200.0) + del f_xyz + + return torch.stack([L, a, b], dim=1) + +def lab_color_transfer( + content_feat: Tensor, + style_feat: Tensor, + luminance_weight: float = 0.8 +) -> Tensor: + content_feat = wavelet_reconstruction(content_feat, style_feat) + + if content_feat.shape != style_feat.shape: + style_feat = F.interpolate( + style_feat, + size=content_feat.shape[-2:], + mode='bilinear', + align_corners=False + ) + + device = content_feat.device + original_dtype = content_feat.dtype + content_feat = content_feat.float() + style_feat = style_feat.float() + + rgb_to_xyz_matrix = torch.tensor([ + [0.4124564, 0.3575761, 0.1804375], + [0.2126729, 0.7151522, 0.0721750], + [0.0193339, 0.1191920, 0.9503041] + ], dtype=torch.float32, device=device) + + xyz_to_rgb_matrix = torch.tensor([ + [ 3.2404542, -1.5371385, -0.4985314], + [-0.9692660, 1.8760108, 0.0415560], + [ 0.0556434, -0.2040259, 1.0572252] + ], dtype=torch.float32, device=device) + + epsilon = CIELAB_DELTA + kappa = CIELAB_KAPPA + + content_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0) + style_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0) + + content_lab = _rgb_to_lab_batch(content_feat, rgb_to_xyz_matrix, epsilon, kappa) + del content_feat + + style_lab = _rgb_to_lab_batch(style_feat, rgb_to_xyz_matrix, epsilon, kappa) + del style_feat, rgb_to_xyz_matrix + + matched_a = _histogram_matching_channel(content_lab[:, 1], style_lab[:, 1]) + matched_b = _histogram_matching_channel(content_lab[:, 2], style_lab[:, 2]) + + if luminance_weight < 1.0: + matched_L = _histogram_matching_channel(content_lab[:, 0], style_lab[:, 0]) + result_L = content_lab[:, 0].mul(luminance_weight).add_(matched_L.mul(1.0 - luminance_weight)) + del matched_L + else: + result_L = content_lab[:, 0] + + del content_lab, style_lab + + result_lab = torch.stack([result_L, matched_a, matched_b], dim=1) + del result_L, matched_a, matched_b + + result_rgb = _lab_to_rgb_batch(result_lab, xyz_to_rgb_matrix, epsilon, kappa) + del result_lab, xyz_to_rgb_matrix + + result = result_rgb.mul_(2.0).sub_(1.0) + del result_rgb + + result = result.to(original_dtype) + + return result + + +def wavelet_color_transfer(content_feat: Tensor, style_feat: Tensor) -> Tensor: + return wavelet_reconstruction(content_feat, style_feat) + + +def adain_color_transfer(content_feat: Tensor, style_feat: Tensor, eps: float = 1e-5) -> Tensor: + if content_feat.shape != style_feat.shape: + style_feat = F.interpolate( + style_feat, + size=content_feat.shape[-2:], + mode='bilinear', + align_corners=False, + ) + + original_dtype = content_feat.dtype + content_feat = content_feat.float() + style_feat = style_feat.float() + + b, c = content_feat.shape[:2] + content_flat = content_feat.reshape(b, c, -1) + style_flat = style_feat.reshape(b, c, -1) + + content_mean = content_flat.mean(dim=2).reshape(b, c, 1, 1) + content_std = (content_flat.var(dim=2, correction=0) + eps).sqrt().reshape(b, c, 1, 1) + style_mean = style_flat.mean(dim=2).reshape(b, c, 1, 1) + style_std = (style_flat.var(dim=2, correction=0) + eps).sqrt().reshape(b, c, 1, 1) + del content_flat, style_flat + + normalized = (content_feat - content_mean) / content_std + del content_mean, content_std + result = normalized * style_std + style_mean + del normalized, style_mean, style_std + + result = result.clamp_(-1.0, 1.0) + if result.dtype != original_dtype: + result = result.to(original_dtype) + return result diff --git a/comfy/ldm/seedvr/constants.py b/comfy/ldm/seedvr/constants.py new file mode 100644 index 000000000..12c4b4bef --- /dev/null +++ b/comfy/ldm/seedvr/constants.py @@ -0,0 +1,48 @@ +"""SeedVR2 constants.""" + +# Temporal chunk-size law: the sampler's activation wall is linear in +# T_latent * pixel area (17-cell resolution sweep + T bisection, RTX 5090, 3b fp16): +# max_latent_frames = (free_GiB - RESERVED - K*SIGMA) / (GIB_PER_MPX_FRAME * megapixels) +# RESERVED covers model staging plus fixed CUDA/torch overhead; SIGMA is the measured +# run-to-run spread of the wall; K=4 trades ~10% smaller chunks for ~1e-5 OOM odds. +SEEDVR2_CHUNK_GIB_PER_MPX_FRAME = 0.55 +SEEDVR2_CHUNK_RESERVED_GIB = 8.5 +SEEDVR2_CHUNK_SIGMA_GIB = 0.55 +SEEDVR2_CHUNK_SIGMA_K = 4 + +SEEDVR2_7B_VID_DIM = 3072 +SEEDVR2_OOM_BACKOFF_DIVISOR = 2 +SEEDVR2_DTYPE_BYTES_FLOOR = 4 +SEEDVR2_7B_MLP_CHUNK = 8192 +SEEDVR2_ROPE_PARTIAL_CHUNK_TOKENS = 4096 # partial-RoPE application token-chunk. +SEEDVR2_LATENT_CHANNELS = 16 + +SEEDVR2_COLOR_MEM_HEADROOM = 0.75 +SEEDVR2_LAB_SCALE_MULTIPLIER = 13 +SEEDVR2_WAVELET_SCALE_MULTIPLIER = 10 # per-frame byte multiplier, wavelet path. +SEEDVR2_ADAIN_SCALE_MULTIPLIER = 6 + +BYTEDANCE_VAE_SCALING_FACTOR = 0.9152 # configs_3b/main.yaml:57. +BYTEDANCE_VAE_SHIFTING_FACTOR = 0.0 +BYTEDANCE_VAE_CONV_MEM_GIB = 0.5 +BYTEDANCE_VAE_NORM_MEM_GIB = 0.5 +BYTEDANCE_LOGVAR_CLAMP_MIN = -30.0 # video_vae_v3/modules/types.py:28. +BYTEDANCE_LOGVAR_CLAMP_MAX = 20.0 # video_vae_v3/modules/types.py:28. +BYTEDANCE_GN_CHUNKS_FP16 = 4 # causal_inflation_lib.py:351 (GroupNorm chunk count, fp16). +BYTEDANCE_GN_CHUNKS_FP32 = 2 # causal_inflation_lib.py:351 (GroupNorm chunk count, fp32). +BYTEDANCE_BLOCK_OUT_CHANNELS = (128, 256, 512, 512) # s8_c16_t4_inflation_sd3.yaml:7-11. +BYTEDANCE_SLICING_SAMPLE_MIN = 4 # s8_c16_t4_inflation_sd3.yaml:22 (slicing_sample_min_size). +BYTEDANCE_VAE_TEMPORAL_DOWNSAMPLE = 4 # infer.py:230 (temporal_downsample_factor); the 4n+1 factor. +BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE = 8 # infer.py:231 (spatial_downsample_factor). +BYTEDANCE_720P_REF_AREA = 45 * 80 # dit_v2/window.py:32 (720p reference area for window scaling). +BYTEDANCE_MAX_TEMPORAL_WINDOW = 30 # dit_v2/window.py:35 (max temporal window frames). +BYTEDANCE_ROPE_MAX_FREQ = 256 # dit_v2/rope.py:31 (pixel-RoPE max frequency). +BYTEDANCE_SINUSOIDAL_DIM = 256 # dit_3b/nadit.py:120 (timestep sinusoidal embed dim). + +ROPE_THETA = 10000 # RoPE base; Su et al., "RoFormer", arXiv:2104.09864. + +CIELAB_DELTA = 6.0 / 29.0 # CIE 15 (delta). +CIELAB_KAPPA = (29.0 / 3.0) ** 3 # CIE 15 (kappa). +D65_WHITE_X = 0.95047 # CIE D65 standard illuminant Xn (Yn = 1). +D65_WHITE_Z = 1.08883 # CIE D65 standard illuminant Zn. +WAVELET_DECOMP_LEVELS = 5 # wavelet color-fix decomposition depth (GIMP/Krita; StableSR). diff --git a/comfy/ldm/seedvr/model.py b/comfy/ldm/seedvr/model.py new file mode 100644 index 000000000..a978698d5 --- /dev/null +++ b/comfy/ldm/seedvr/model.py @@ -0,0 +1,1361 @@ +from dataclasses import dataclass +from typing import Optional, Tuple, Union, List, Dict, Any, Callable +import torch.nn.functional as F +from math import ceil, pi +import torch +from itertools import accumulate, chain +from comfy.ldm.modules.diffusionmodules.model import get_timestep_embedding +from comfy.ldm.seedvr.attention import optimized_var_attention +from torch.nn.modules.utils import _triple +from torch import nn +import math +from comfy.ldm.flux.math import apply_rope1 +from comfy.ldm.seedvr.constants import ( + BYTEDANCE_720P_REF_AREA, + BYTEDANCE_MAX_TEMPORAL_WINDOW, + BYTEDANCE_ROPE_MAX_FREQ, + BYTEDANCE_SINUSOIDAL_DIM, + ROPE_THETA, + SEEDVR2_7B_MLP_CHUNK, + SEEDVR2_7B_VID_DIM, + SEEDVR2_LATENT_CHANNELS, + SEEDVR2_ROPE_PARTIAL_CHUNK_TOKENS, +) +import comfy.model_management +import comfy.ops + +class Cache: + def __init__(self, disable=False, prefix="", cache=None): + self.cache = cache if cache is not None else {} + self.disable = disable + self.prefix = prefix + + def __call__(self, key: str, fn: Callable): + if self.disable: + return fn() + + key = self.prefix + key + if key not in self.cache: + result = fn() + self.cache[key] = result + return self.cache[key] + + def namespace(self, namespace: str): + return Cache( + disable=self.disable, + prefix=self.prefix + namespace + ".", + cache=self.cache, + ) + +def repeat_concat( + vid: torch.FloatTensor, # (VL ... c) + txt: torch.FloatTensor, # (TL ... c) + vid_len: torch.LongTensor, # (n*b) + txt_len: torch.LongTensor, # (b) + txt_repeat: List, # (n) +) -> torch.FloatTensor: # (L ... c) + vid = torch.split(vid, vid_len.tolist()) + txt = torch.split(txt, txt_len.tolist()) + txt = [[x] * n for x, n in zip(txt, txt_repeat)] + txt = list(chain(*txt)) + return torch.cat(list(chain(*zip(vid, txt)))) + +def repeat_concat_idx( + vid_len: torch.LongTensor, # (n*b) + txt_len: torch.LongTensor, # (b) + txt_repeat: torch.LongTensor, # (n) +) -> Tuple[ + Callable, + Callable, +]: + device = vid_len.device + vid_idx = torch.arange(vid_len.sum(), device=device) + txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device) + txt_repeat_list = txt_repeat.tolist() + tgt_idx = repeat_concat(vid_idx, txt_idx, vid_len, txt_len, txt_repeat_list) + src_idx = torch.argsort(tgt_idx) + txt_idx_len = len(tgt_idx) - len(vid_idx) + repeat_txt_len = (txt_len * txt_repeat).tolist() + + def unconcat_coalesce(all): + vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len]) + txt_out_coalesced = [] + for txt, repeat_time in zip(txt_out.split(repeat_txt_len), txt_repeat_list): + txt = txt.reshape(-1, repeat_time, *txt.shape[1:]).mean(1) + txt_out_coalesced.append(txt) + return vid_out, torch.cat(txt_out_coalesced) + + return ( + lambda vid, txt: torch.cat([vid, txt])[tgt_idx], + lambda all: unconcat_coalesce(all), + ) + +def cumulative_lengths(lengths): + return [0, *accumulate(lengths)] + + +@dataclass +class MMArg: + vid: Any + txt: Any + +def get_args(key: str, args: List[Any]) -> List[Any]: + return [getattr(v, key) if isinstance(v, MMArg) else v for v in args] + + +def get_kwargs(key: str, kwargs: Dict[str, Any]) -> Dict[str, Any]: + return {k: getattr(v, key) if isinstance(v, MMArg) else v for k, v in kwargs.items()} + + +def get_window_op(name: str): + if name == "720pwin_by_size_bysize": + return make_720Pwindows_bysize + if name == "720pswin_by_size_bysize": + return make_shifted_720Pwindows_bysize + raise ValueError(f"Unknown windowing method: {name}") + + +def make_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]): + t, h, w = size + resized_nt, resized_nh, resized_nw = num_windows + scale = math.sqrt(BYTEDANCE_720P_REF_AREA / (h * w)) + resized_h, resized_w = round(h * scale), round(w * scale) + wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) + wt = ceil(min(t, BYTEDANCE_MAX_TEMPORAL_WINDOW) / resized_nt) + nt, nh, nw = ceil(t / wt), ceil(h / wh), ceil(w / ww) + return [ + ( + slice(it * wt, min((it + 1) * wt, t)), + slice(ih * wh, min((ih + 1) * wh, h)), + slice(iw * ww, min((iw + 1) * ww, w)), + ) + for iw in range(nw) + if min((iw + 1) * ww, w) > iw * ww + for ih in range(nh) + if min((ih + 1) * wh, h) > ih * wh + for it in range(nt) + if min((it + 1) * wt, t) > it * wt + ] + +def make_shifted_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]): + t, h, w = size + resized_nt, resized_nh, resized_nw = num_windows + scale = math.sqrt(BYTEDANCE_720P_REF_AREA / (h * w)) + resized_h, resized_w = round(h * scale), round(w * scale) + wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) + wt = ceil(min(t, BYTEDANCE_MAX_TEMPORAL_WINDOW) / resized_nt) + + st, sh, sw = ( + 0.5 if wt < t else 0, + 0.5 if wh < h else 0, + 0.5 if ww < w else 0, + ) + nt, nh, nw = ceil((t - st) / wt), ceil((h - sh) / wh), ceil((w - sw) / ww) + nt, nh, nw = ( + nt + 1 if st > 0 else 1, + nh + 1 if sh > 0 else 1, + nw + 1 if sw > 0 else 1, + ) + return [ + ( + slice(max(int((it - st) * wt), 0), min(int((it - st + 1) * wt), t)), + slice(max(int((ih - sh) * wh), 0), min(int((ih - sh + 1) * wh), h)), + slice(max(int((iw - sw) * ww), 0), min(int((iw - sw + 1) * ww), w)), + ) + for iw in range(nw) + if min(int((iw - sw + 1) * ww), w) > max(int((iw - sw) * ww), 0) + for ih in range(nh) + if min(int((ih - sh + 1) * wh), h) > max(int((ih - sh) * wh), 0) + for it in range(nt) + if min(int((it - st + 1) * wt), t) > max(int((it - st) * wt), 0) + ] + +class RotaryEmbedding(nn.Module): + def __init__( + self, + dim, + freqs_for = 'lang', + theta = 10000, + max_freq = 10, + ): + super().__init__() + + self.freqs_for = freqs_for + + if freqs_for == 'lang': + freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) + elif freqs_for == 'pixel': + freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi + else: + raise ValueError(f"Unknown rotary frequency type: {freqs_for}") + + self.register_buffer("freqs", freqs) + + @property + def device(self): + return self.freqs.device + + def get_axial_freqs( + self, + *dims, + offsets = None + ): + Colon = slice(None) + all_freqs = [] + + if exists(offsets): + if len(offsets) != len(dims): + raise ValueError(f"SeedVR2 rotary offsets length must match dims length, got {len(offsets)} and {len(dims)}.") + + for ind, dim in enumerate(dims): + + offset = 0 + if exists(offsets): + offset = offsets[ind] + + if self.freqs_for == 'pixel': + pos = torch.linspace(-1, 1, steps = dim, device = self.device) + else: + pos = torch.arange(dim, device = self.device) + + pos = pos + offset + + freqs = self.forward(pos) + + all_axis = [None] * len(dims) + all_axis[ind] = Colon + + new_axis_slice = (Ellipsis, *all_axis, Colon) + all_freqs.append(freqs[new_axis_slice]) + + all_freqs = torch.broadcast_tensors(*all_freqs) + return torch.cat(all_freqs, dim = -1) + + def forward( + self, + t, + ): + freqs = self.freqs + + freqs = torch.einsum('..., f -> ... f', t.type(freqs.dtype), freqs) + freqs = freqs.unsqueeze(-1).expand(*freqs.shape, 2).flatten(-2) + + return freqs + +class RotaryEmbeddingBase(nn.Module): + def __init__(self, dim: int, rope_dim: int): + super().__init__() + self.rope = RotaryEmbedding( + dim=dim // rope_dim, + freqs_for="pixel", + max_freq=BYTEDANCE_ROPE_MAX_FREQ, + ) + + def get_axial_freqs(self, *dims): + return self.rope.get_axial_freqs(*dims) + + +class RotaryEmbedding3d(RotaryEmbeddingBase): + def __init__(self, dim: int): + super().__init__(dim, rope_dim=3) + self.mm = False + + +class NaRotaryEmbedding3d(RotaryEmbedding3d): + def forward( + self, + q: torch.FloatTensor, + k: torch.FloatTensor, + shape: torch.LongTensor, + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + freqs = cache("rope_freqs_3d", lambda: self.get_freqs(shape)) + freqs = freqs.to(device=q.device) + q = q.transpose(0, 1) + k = k.transpose(0, 1) + q = _apply_seedvr2_rotary_emb(freqs, q.float()).to(q.dtype) + k = _apply_seedvr2_rotary_emb(freqs, k.float()).to(k.dtype) + q = q.transpose(0, 1) + k = k.transpose(0, 1) + return q, k + + @torch._dynamo.disable + def get_freqs( + self, + shape: torch.LongTensor, + ) -> torch.Tensor: + # Primary provenance: ByteDance-Seed/SeedVR models/dit/rope.py builds + # 7B pixel RoPE with the interleaved-angle convention, not Comfy's + # Flux freqs_cis matrix. + plain_rope = RotaryEmbedding( + dim=self.rope.freqs.numel() * 2, + freqs_for="pixel", + max_freq=BYTEDANCE_ROPE_MAX_FREQ, + ) + plain_rope = plain_rope.to(self.rope.device) + freq_list = [] + for f, h, w in shape.tolist(): + freqs = plain_rope.get_axial_freqs(f, h, w) + freq_list.append(freqs.view(-1, freqs.size(-1))) + return torch.cat(freq_list, dim=0) + + +class MMRotaryEmbeddingBase(RotaryEmbeddingBase): + def __init__(self, dim: int, rope_dim: int): + super().__init__(dim, rope_dim) + self.rope = RotaryEmbedding( + dim=dim // rope_dim, + freqs_for="lang", + theta=ROPE_THETA, + ) + self.mm = True + +def slice_at_dim(t, dim_slice: slice, *, dim): + dim += (t.ndim if dim < 0 else 0) + colons = [slice(None)] * t.ndim + colons[dim] = dim_slice + return t[tuple(colons)] + +def rotate_half(x): + x = x.reshape(*x.shape[:-1], x.shape[-1] // 2, 2) + x1, x2 = x.unbind(dim = -1) + x = torch.stack((-x2, x1), dim = -1) + return x.flatten(-2) +def exists(val): + return val is not None + +def _apply_seedvr2_rotary_emb( + freqs: torch.Tensor, + t: torch.Tensor, + start_index: int = 0, + scale: float = 1.0, + seq_dim: int = -2, + freqs_seq_dim: int | None = None, +) -> torch.Tensor: + dtype = t.dtype + if freqs_seq_dim is None and (freqs.ndim == 2 or t.ndim == 3): + freqs_seq_dim = 0 + + if t.ndim == 3 or freqs_seq_dim is not None: + seq_len = t.shape[seq_dim] + freqs = slice_at_dim(freqs, slice(-seq_len, None), dim=freqs_seq_dim) + + rot_feats = freqs.shape[-1] + end_index = start_index + rot_feats + + t_left = t[..., :start_index] + t_middle = t[..., start_index:end_index] + t_right = t[..., end_index:] + + freqs = freqs.to(device=t_middle.device, dtype=t_middle.dtype) + cos = freqs.cos() * scale + sin = freqs.sin() * scale + t_middle = (t_middle * cos) + (rotate_half(t_middle) * sin) + return torch.cat((t_left, t_middle, t_right), dim=-1).to(dtype) + +def _to_flux_freqs_cis(freqs_interleaved: torch.Tensor) -> torch.Tensor: + angles = freqs_interleaved[..., ::2].float() + cos = torch.cos(angles) + sin = torch.sin(angles) + out = torch.stack([cos, -sin, sin, cos], dim=-1) + return out.reshape(*out.shape[:-1], 2, 2) + + +def _apply_rope1_partial(t: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + out = t.clone() if t.requires_grad or comfy.model_management.in_training else t + rot_d = 2 * freqs_cis.shape[-3] + seq_len = out.shape[-2] + for start in range(0, seq_len, SEEDVR2_ROPE_PARTIAL_CHUNK_TOKENS): + end = min(start + SEEDVR2_ROPE_PARTIAL_CHUNK_TOKENS, seq_len) + freqs_chunk = freqs_cis[start:end] + if rot_d == out.shape[-1]: + out[..., start:end, :] = apply_rope1(out[..., start:end, :], freqs_chunk).to(out.dtype) + else: + out[..., start:end, :rot_d] = apply_rope1(out[..., start:end, :rot_d], freqs_chunk).to(out.dtype) + return out + + +class NaMMRotaryEmbedding3d(MMRotaryEmbeddingBase): + def __init__(self, dim: int): + super().__init__(dim, rope_dim=3) + + def forward( + self, + vid_q: torch.FloatTensor, # L h d + vid_k: torch.FloatTensor, # L h d + vid_shape: torch.LongTensor, # B 3 + txt_q: torch.FloatTensor, # L h d + txt_k: torch.FloatTensor, # L h d + txt_shape: torch.LongTensor, # B 1 + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + torch.FloatTensor, + torch.FloatTensor, + ]: + vid_freqs, txt_freqs = cache( + "mmrope_freqs_3d", + lambda: self.get_freqs(vid_shape, txt_shape), + ) + target_device = vid_q.device + if vid_freqs.device != target_device: + vid_freqs = vid_freqs.to(target_device) + if txt_freqs.device != target_device: + txt_freqs = txt_freqs.to(target_device) + vid_q = vid_q.transpose(0, 1) + vid_k = vid_k.transpose(0, 1) + vid_q = _apply_rope1_partial(vid_q, vid_freqs) + vid_k = _apply_rope1_partial(vid_k, vid_freqs) + vid_q = vid_q.transpose(0, 1) + vid_k = vid_k.transpose(0, 1) + + txt_q = txt_q.transpose(0, 1) + txt_k = txt_k.transpose(0, 1) + txt_q = _apply_rope1_partial(txt_q, txt_freqs) + txt_k = _apply_rope1_partial(txt_k, txt_freqs) + txt_q = txt_q.transpose(0, 1) + txt_k = txt_k.transpose(0, 1) + return vid_q, vid_k, txt_q, txt_k + + @torch._dynamo.disable # Disable compilation: .tolist() is data-dependent and causes graph breaks + def get_freqs( + self, + vid_shape: torch.LongTensor, + txt_shape: torch.LongTensor, + ) -> Tuple[ + torch.Tensor, + torch.Tensor, + ]: + + max_temporal = 0 + max_height = 0 + max_width = 0 + max_txt_len = 0 + + for (f, h, w), l in zip(vid_shape.tolist(), txt_shape[:, 0].tolist()): + max_temporal = max(max_temporal, l + f) + max_height = max(max_height, h) + max_width = max(max_width, w) + max_txt_len = max(max_txt_len, l) + + autocast_device = "cuda" if torch.cuda.is_available() else "cpu" + with torch.amp.autocast(autocast_device, enabled=False): + vid_freqs = self.get_axial_freqs( + max_temporal + 16, + max_height + 4, + max_width + 4, + ).float() + txt_freqs = self.get_axial_freqs(max_txt_len + 16) + + vid_freq_list, txt_freq_list = [], [] + for (f, h, w), l in zip(vid_shape.tolist(), txt_shape[:, 0].tolist()): + vid_freq = vid_freqs[l : l + f, :h, :w].reshape(-1, vid_freqs.size(-1)) + txt_freq = txt_freqs[:l].repeat(1, 3).reshape(-1, vid_freqs.size(-1)) + vid_freq_list.append(vid_freq) + txt_freq_list.append(txt_freq) + vid_freqs_interleaved = torch.cat(vid_freq_list, dim=0) + txt_freqs_interleaved = torch.cat(txt_freq_list, dim=0) + + return _to_flux_freqs_cis(vid_freqs_interleaved), _to_flux_freqs_cis(txt_freqs_interleaved) + +class MMModule(nn.Module): + def __init__( + self, + module: Callable[..., nn.Module], + *args, + shared_weights: bool = False, + vid_only: bool = False, + **kwargs, + ): + super().__init__() + self.shared_weights = shared_weights + self.vid_only = vid_only + if self.shared_weights: + if get_args("vid", args) != get_args("txt", args): + raise ValueError("SeedVR2 shared MMModule requires matching vid/txt args.") + if get_kwargs("vid", kwargs) != get_kwargs("txt", kwargs): + raise ValueError("SeedVR2 shared MMModule requires matching vid/txt kwargs.") + self.all = module(*get_args("vid", args), **get_kwargs("vid", kwargs)) + else: + self.vid = module(*get_args("vid", args), **get_kwargs("vid", kwargs)) + self.txt = ( + module(*get_args("txt", args), **get_kwargs("txt", kwargs)) + if not vid_only + else None + ) + + def forward( + self, + vid: torch.FloatTensor, + txt: torch.FloatTensor, + *args, + **kwargs, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + vid_module = self.vid if not self.shared_weights else self.all + vid = vid_module(vid, *get_args("vid", args), **get_kwargs("vid", kwargs)) + if not self.vid_only: + txt_module = self.txt if not self.shared_weights else self.all + txt = txt.to(device=vid.device, dtype=vid.dtype) + txt = txt_module(txt, *get_args("txt", args), **get_kwargs("txt", kwargs)) + return vid, txt + +def get_na_rope(rope_type: Optional[str], dim: int): + if rope_type is None: + return None + if rope_type == "rope3d": + return NaRotaryEmbedding3d(dim=dim) + if rope_type == "mmrope3d": + return NaMMRotaryEmbedding3d(dim=dim) + raise ValueError(f"Unknown SeedVR2 rope type: {rope_type}") + +class NaMMAttention(nn.Module): + def __init__( + self, + vid_dim: int, + txt_dim: int, + heads: int, + head_dim: int, + qk_bias: bool, + qk_norm, + qk_norm_eps: float, + rope_type: Optional[str], + rope_dim: int, + shared_weights: bool, + device, dtype, operations, + ): + super().__init__() + dim = MMArg(vid_dim, txt_dim) + self.heads = heads + inner_dim = heads * head_dim + qkv_dim = inner_dim * 3 + self.head_dim = head_dim + self.proj_qkv = MMModule( + operations.Linear, dim, qkv_dim, bias=qk_bias, shared_weights=shared_weights, device=device, dtype=dtype + ) + self.proj_out = MMModule(operations.Linear, inner_dim, dim, shared_weights=shared_weights, device=device, dtype=dtype) + self.norm_q = MMModule( + qk_norm, + normalized_shape=head_dim, + eps=qk_norm_eps, + elementwise_affine=True, + shared_weights=shared_weights, + device=device, dtype=dtype + ) + self.norm_k = MMModule( + qk_norm, + normalized_shape=head_dim, + eps=qk_norm_eps, + elementwise_affine=True, + shared_weights=shared_weights, + device=device, dtype=dtype + ) + + + self.rope = get_na_rope(rope_type=rope_type, dim=rope_dim) + +def window( + hid: torch.FloatTensor, # (L c) + hid_shape: torch.LongTensor, # (b n) + window_fn: Callable[[torch.Tensor], List[torch.Tensor]], +): + hid = unflatten(hid, hid_shape) + hid = list(map(window_fn, hid)) + hid_windows_list = [len(x) for x in hid] + hid_windows = torch.as_tensor(hid_windows_list, device=hid_shape.device) + hid = list(chain(*hid)) + hid_len_list = [math.prod(x.shape[:-1]) for x in hid] + hid, hid_shape = flatten(hid) + return hid, hid_shape, hid_windows, hid_len_list, hid_windows_list + +def window_idx( + hid_shape: torch.LongTensor, # (b n) + window_fn: Callable[[torch.Tensor], List[torch.Tensor]], +): + hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1) + tgt_idx, tgt_shape, tgt_windows, tgt_len_list, tgt_windows_list = window(hid_idx, hid_shape, window_fn) + tgt_idx = tgt_idx.squeeze(-1) + src_idx = torch.argsort(tgt_idx) + return ( + lambda hid: torch.index_select(hid, 0, tgt_idx), + lambda hid: torch.index_select(hid, 0, src_idx), + tgt_shape, + tgt_windows, + tgt_len_list, + tgt_windows_list, + ) + +class NaSwinAttention(NaMMAttention): + def __init__( + self, + *args, + window: Union[int, Tuple[int, int, int]], + window_method: str, + version: bool = False, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.version_7b = version + self.window = _triple(window) + self.window_method = window_method + if not all(isinstance(v, int) and v >= 0 for v in self.window): + raise ValueError(f"SeedVR2 window must contain non-negative integers, got {self.window}.") + + self.window_op = get_window_op(window_method) + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + + vid_qkv, txt_qkv = self.proj_qkv(vid, txt) + + cache_win = cache.namespace(f"{self.window_method}_{self.window}_sd3") + + def make_window(x: torch.Tensor): + t, h, w, _ = x.shape + window_slices = self.window_op((t, h, w), self.window) + return [x[st, sh, sw] for (st, sh, sw) in window_slices] + + window_partition, window_reverse, window_shape, window_count, vid_len_win_list, window_count_list = cache_win( + "win_transform", + lambda: window_idx(vid_shape, make_window), + ) + vid_qkv_win = window_partition(vid_qkv) + + vid_qkv_win = vid_qkv_win.reshape(vid_qkv_win.shape[0], 3, self.heads, self.head_dim) + txt_qkv = txt_qkv.reshape(txt_qkv.shape[0], 3, self.heads, self.head_dim) + + vid_q, vid_k, vid_v = vid_qkv_win.unbind(1) + txt_q, txt_k, txt_v = txt_qkv.unbind(1) + + vid_q, txt_q = self.norm_q(vid_q, txt_q) + vid_k, txt_k = self.norm_k(vid_k, txt_k) + + txt_len = cache("txt_len", lambda: txt_shape.prod(-1)) + + vid_len_win = cache_win("vid_len", lambda: window_shape.prod(-1)) + txt_len = txt_len.to(window_count.device) + + if self.rope: + if self.version_7b: + vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win) + elif self.rope.mm: + _, num_h, _ = txt_q.shape + txt_q_repeat = txt_q.flatten(1, 2) + txt_q_repeat = unflatten(txt_q_repeat, txt_shape) + txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count_list)] + txt_q_repeat = list(chain(*txt_q_repeat)) + txt_q_repeat, txt_shape_repeat = flatten(txt_q_repeat) + txt_q_repeat = txt_q_repeat.reshape(txt_q_repeat.shape[0], num_h, self.head_dim) + + txt_k_repeat = txt_k.flatten(1, 2) + txt_k_repeat = unflatten(txt_k_repeat, txt_shape) + txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count_list)] + txt_k_repeat = list(chain(*txt_k_repeat)) + txt_k_repeat, _ = flatten(txt_k_repeat) + txt_k_repeat = txt_k_repeat.reshape(txt_k_repeat.shape[0], num_h, self.head_dim) + + vid_q, vid_k, txt_q, txt_k = self.rope( + vid_q, vid_k, window_shape, txt_q_repeat, txt_k_repeat, txt_shape_repeat, cache_win + ) + else: + vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win) + + txt_len_win_list = cache_win( + "txt_len_list", + lambda: [txt_len for txt_len, window_count in zip(txt_len.tolist(), window_count_list) for _ in range(window_count)], + ) + all_len_win = cache_win("all_len", lambda: [vid_len + txt_len for vid_len, txt_len in zip(vid_len_win_list, txt_len_win_list)]) + concat_win, unconcat_win = cache_win( + "mm_pnp", lambda: repeat_concat_idx(vid_len_win, txt_len, window_count) + ) + out = optimized_var_attention( + q=concat_win(vid_q, txt_q), + k=concat_win(vid_k, txt_k), + v=concat_win(vid_v, txt_v), + heads=self.heads, skip_reshape=True, skip_output_reshape=True, + cu_seqlens_q=cache_win("vid_seqlens_q", lambda: cumulative_lengths(all_len_win)), + cu_seqlens_k=cache_win("vid_seqlens_k", lambda: cumulative_lengths(all_len_win)), + ) + vid_out, txt_out = unconcat_win(out) + + vid_out = vid_out.flatten(1, 2) + txt_out = txt_out.flatten(1, 2) + vid_out = window_reverse(vid_out) + + vid_out, txt_out = self.proj_out(vid_out, txt_out) + + return vid_out, txt_out + +class MLP(nn.Module): + def __init__( + self, + dim: int, + expand_ratio: int, + device, dtype, operations + ): + super().__init__() + self.proj_in = operations.Linear(dim, dim * expand_ratio, device=device, dtype=dtype) + self.act = nn.GELU("tanh") + self.proj_out = operations.Linear(dim * expand_ratio, dim, device=device, dtype=dtype) + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + x = self.proj_in(x) + x = self.act(x) + x = self.proj_out(x) + return x + + +class SwiGLUMLP(nn.Module): + def __init__( + self, + dim: int, + expand_ratio: int, + multiple_of: int = 256, + device=None, dtype=None, operations=None + ): + super().__init__() + hidden_dim = int(2 * dim * expand_ratio / 3) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + self.proj_in_gate = operations.Linear(dim, hidden_dim, bias=False, device=device, dtype=dtype) + self.proj_out = operations.Linear(hidden_dim, dim, bias=False, device=device, dtype=dtype) + self.proj_in = operations.Linear(dim, hidden_dim, bias=False, device=device, dtype=dtype) + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + return self.proj_out(F.silu(self.proj_in_gate(x)) * self.proj_in(x)) + +def get_mlp(mlp_type: Optional[str] = "normal"): + if mlp_type == "normal": + return MLP + if mlp_type == "swiglu": + return SwiGLUMLP + raise ValueError(f"Unknown SeedVR2 MLP type: {mlp_type}") + +class NaMMSRTransformerBlock(nn.Module): + def __init__( + self, + *, + vid_dim: int, + txt_dim: int, + emb_dim: int, + heads: int, + head_dim: int, + expand_ratio: int, + norm, + norm_eps: float, + ada, + qk_bias: bool, + qk_norm, + mlp_type: str, + shared_weights: bool, + rope_type: str, + rope_dim: int, + is_last_layer: bool, + window: Union[int, Tuple[int, int, int]], + window_method: str, + version: bool, + device, dtype, operations, + ): + super().__init__() + dim = MMArg(vid_dim, txt_dim) + self.attn_norm = MMModule(norm, normalized_shape=dim, eps=norm_eps, elementwise_affine=False, shared_weights=shared_weights, device=device, dtype=dtype) + + self.attn = NaSwinAttention( + vid_dim=vid_dim, + txt_dim=txt_dim, + heads=heads, + head_dim=head_dim, + qk_bias=qk_bias, + qk_norm=qk_norm, + qk_norm_eps=norm_eps, + rope_type=rope_type, + rope_dim=rope_dim, + shared_weights=shared_weights, + window=window, + window_method=window_method, + version=version, + device=device, dtype=dtype, operations=operations + ) + + self.mlp_norm = MMModule(norm, normalized_shape=dim, eps=norm_eps, elementwise_affine=False, shared_weights=shared_weights, vid_only=is_last_layer, device=device, dtype=dtype) + self.mlp = MMModule( + get_mlp(mlp_type), + dim=dim, + expand_ratio=expand_ratio, + shared_weights=shared_weights, + vid_only=is_last_layer, + device=device, dtype=dtype, operations=operations + ) + self.ada = MMModule(ada, dim=dim, emb_dim=emb_dim, layers=["attn", "mlp"], shared_weights=shared_weights, vid_only=is_last_layer, device=device, dtype=dtype) + self.is_last_layer = is_last_layer + self.version = version + + def _seedvr2_7b_mlp( + self, + vid: torch.FloatTensor, + txt: torch.FloatTensor, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + vid_module = self.mlp.vid if not self.mlp.shared_weights else self.mlp.all + if comfy.model_management.in_training or vid.requires_grad: + vid = torch.cat([vid_module(chunk) for chunk in vid.split(SEEDVR2_7B_MLP_CHUNK, dim=0)], dim=0) + else: + vid_out = None + offset = 0 + for chunk in vid.split(SEEDVR2_7B_MLP_CHUNK, dim=0): + chunk_out = vid_module(chunk) + if vid_out is None: + vid_out = chunk_out.new_empty((vid.shape[0], *chunk_out.shape[1:])) + vid_out[offset:offset + chunk_out.shape[0]] = chunk_out + offset += chunk_out.shape[0] + vid = vid_out + if not self.mlp.vid_only: + txt_module = self.mlp.txt if not self.mlp.shared_weights else self.mlp.all + txt = txt.to(device=vid.device, dtype=vid.dtype) + txt = txt_module(txt) + return vid, txt + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + emb: torch.FloatTensor, + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + torch.LongTensor, + torch.LongTensor, + ]: + hid_len = MMArg( + cache("vid_len", lambda: vid_shape.prod(-1)), + cache("txt_len", lambda: txt_shape.prod(-1)), + ) + ada_kwargs = { + "emb": emb, + "hid_len": hid_len, + "cache": cache, + "branch_tag": MMArg("vid", "txt"), + } + + vid_attn, txt_attn = self.attn_norm(vid, txt) + vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="in", **ada_kwargs) + vid_attn, txt_attn = self.attn(vid_attn, txt_attn, vid_shape, txt_shape, cache) + vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="out", **ada_kwargs) + vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt) + + vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn) + vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="in", **ada_kwargs) + if self.version: + vid_mlp, txt_mlp = self._seedvr2_7b_mlp(vid_mlp, txt_mlp) + else: + vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp) + vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="out", **ada_kwargs) + vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn) + + return vid_mlp, txt_mlp, vid_shape, txt_shape + +class PatchOut(nn.Module): + def __init__( + self, + out_channels: int, + patch_size: Union[int, Tuple[int, int, int]], + dim: int, + device, dtype, operations + ): + super().__init__() + t, h, w = _triple(patch_size) + self.patch_size = t, h, w + self.proj = operations.Linear(dim, out_channels * t * h * w, device=device, dtype=dtype) + + def forward( + self, + vid: torch.Tensor, + ) -> torch.Tensor: + t, h, w = self.patch_size + vid = self.proj(vid) + b, T, H, W, channels = vid.shape + c = channels // (t * h * w) + vid = vid.view(b, T, H, W, t, h, w, c).permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(b, c, T * t, H * h, W * w) + if t > 1: + vid = vid[:, :, (t - 1) :] + return vid + +class NaPatchOut(PatchOut): + def forward( + self, + vid: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, + cache: Optional[Cache] = None, + vid_shape_before_patchify = None + ) -> Tuple[ + torch.FloatTensor, + torch.LongTensor, + ]: + if cache is None: + cache = Cache(disable=True) + + t, h, w = self.patch_size + vid = self.proj(vid) + + if not (t == h == w == 1): + vid = unflatten(vid, vid_shape) + for i in range(len(vid)): + T, H, W, channels = vid[i].shape + c = channels // (t * h * w) + vid[i] = vid[i].view(T, H, W, t, h, w, c).permute(0, 3, 1, 4, 2, 5, 6).reshape(T * t, H * h, W * w, c) + if t > 1 and vid_shape_before_patchify[i, 0] % t != 0: + vid[i] = vid[i][(t - vid_shape_before_patchify[i, 0] % t) :] + vid, vid_shape = flatten(vid) + + return vid, vid_shape + +class PatchIn(nn.Module): + def __init__( + self, + in_channels: int, + patch_size: Union[int, Tuple[int, int, int]], + dim: int, + device, dtype, operations + ): + super().__init__() + t, h, w = _triple(patch_size) + self.patch_size = t, h, w + self.proj = operations.Linear(in_channels * t * h * w, dim, device=device, dtype=dtype) + + def forward( + self, + vid: torch.Tensor, + ) -> torch.Tensor: + t, h, w = self.patch_size + if t > 1: + if vid.size(2) % t != 1: + raise ValueError( + f"SeedVR2 patch input temporal size must satisfy T % {t} == 1, got {vid.size(2)}." + ) + vid = torch.cat([vid[:, :, :1]] * (t - 1) + [vid], dim=2) + b, c, Tt, Hh, Ww = vid.shape + vid = vid.view(b, c, Tt // t, t, Hh // h, h, Ww // w, w).permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(b, Tt // t, Hh // h, Ww // w, t * h * w * c) + vid = self.proj(vid) + return vid + +class NaPatchIn(PatchIn): + def forward( + self, + vid: torch.Tensor, # l c + vid_shape: torch.LongTensor, + cache: Optional[Cache] = None, + ) -> torch.Tensor: + if cache is None: + cache = Cache(disable=True) + cache = cache.namespace("patch") + vid_shape_before_patchify = cache("vid_shape_before_patchify", lambda: vid_shape) + t, h, w = self.patch_size + if not (t == h == w == 1): + vid = unflatten(vid, vid_shape) + for i in range(len(vid)): + if t > 1 and vid_shape_before_patchify[i, 0] % t != 0: + vid[i] = torch.cat([vid[i][:1]] * (t - vid[i].size(0) % t) + [vid[i]], dim=0) + Tt, Hh, Ww, c = vid[i].shape + vid[i] = vid[i].view(Tt // t, t, Hh // h, h, Ww // w, w, c).permute(0, 2, 4, 1, 3, 5, 6).reshape(Tt // t, Hh // h, Ww // w, t * h * w * c) + vid, vid_shape = flatten(vid) + + vid = self.proj(vid) + return vid, vid_shape + +def expand_dims(x: torch.Tensor, dim: int, ndim: int): + shape = x.shape + shape = shape[:dim] + (1,) * (ndim - len(shape)) + shape[dim:] + return x.reshape(shape) + + +class AdaSingle(nn.Module): + def __init__( + self, + dim: int, + emb_dim: int, + layers: List[str], + modes: Tuple[str, ...] = ("in", "out"), + device = None, dtype = None, + ): + if emb_dim != 6 * dim: + raise ValueError(f"SeedVR2 AdaSingle requires emb_dim == 6 * dim, got emb_dim={emb_dim}, dim={dim}.") + super().__init__() + self.dim = dim + self.emb_dim = emb_dim + self.layers = layers + + param_kwargs = {"device": device, "dtype": dtype} + + for l in layers: + if "in" in modes: + self.register_parameter(f"{l}_shift", nn.Parameter(torch.empty(dim, **param_kwargs))) + self.register_parameter(f"{l}_scale", nn.Parameter(torch.empty(dim, **param_kwargs))) + if "out" in modes: + self.register_parameter(f"{l}_gate", nn.Parameter(torch.empty(dim, **param_kwargs))) + + def forward( + self, + hid: torch.FloatTensor, # b ... c + emb: torch.FloatTensor, # b d + layer: str, + mode: str, + cache: Optional[Cache] = None, + branch_tag: str = "", + hid_len: Optional[torch.LongTensor] = None, # b + ) -> torch.FloatTensor: + if cache is None: + cache = Cache(disable=True) + idx = self.layers.index(layer) + emb = emb.reshape(emb.shape[0], -1, len(self.layers), 3)[:, :, idx, :] + emb = expand_dims(emb, 1, hid.ndim + 1) + + if hid_len is not None: + emb = cache( + f"emb_repeat_{idx}_{branch_tag}", + lambda: torch.repeat_interleave(emb, hid_len, dim=0), + ) + + shiftA, scaleA, gateA = emb.unbind(-1) + shiftB, scaleB, gateB = ( + getattr(self, f"{layer}_shift", None), + getattr(self, f"{layer}_scale", None), + getattr(self, f"{layer}_gate", None), + ) + + if mode == "in": + shiftB = comfy.ops.cast_to_input(shiftB, hid) + scaleB = comfy.ops.cast_to_input(scaleB, hid) + return hid.mul_(scaleA + scaleB).add_(shiftA + shiftB) + if mode == "out": + if gateB is not None: + gateB = comfy.ops.cast_to_input(gateB, hid) + return hid.mul_(gateA + gateB) + else: + return hid.mul_(gateA) + + raise ValueError(f"Unknown AdaSingle mode: {mode}") + + +class TimeEmbedding(nn.Module): + def __init__( + self, + sinusoidal_dim: int, + hidden_dim: int, + output_dim: int, + device, dtype, operations + ): + super().__init__() + self.sinusoidal_dim = sinusoidal_dim + self.proj_in = operations.Linear(sinusoidal_dim, hidden_dim, device=device, dtype=dtype) + self.proj_hid = operations.Linear(hidden_dim, hidden_dim, device=device, dtype=dtype) + self.proj_out = operations.Linear(hidden_dim, output_dim, device=device, dtype=dtype) + self.act = nn.SiLU() + + def forward( + self, + timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], + device: torch.device, + dtype: torch.dtype, + ) -> torch.FloatTensor: + if not torch.is_tensor(timestep): + timestep = torch.tensor([timestep], device=device, dtype=dtype) + if timestep.ndim == 0: + timestep = timestep[None] + + emb = get_timestep_embedding( + timesteps=timestep, + embedding_dim=self.sinusoidal_dim, + flip_sin_to_cos=False, + downscale_freq_shift=0, + ).to(dtype) + emb = self.proj_in(emb) + emb = self.act(emb) + emb = self.proj_hid(emb) + emb = self.act(emb) + emb = self.proj_out(emb) + return emb + +def flatten( + hid: List[torch.FloatTensor], # List of (*** c) +) -> Tuple[ + torch.FloatTensor, # (L c) + torch.LongTensor, # (b n) +]: + if len(hid) == 0: + raise ValueError("SeedVR2 flatten requires at least one tensor.") + shape = torch.as_tensor([x.shape[:-1] for x in hid], device=hid[0].device) + hid = torch.cat([x.flatten(0, -2) for x in hid]) + return hid, shape + + +def unflatten( + hid: torch.FloatTensor, # (L c) or (L ... c) + hid_shape: torch.LongTensor, # (b n) +) -> List[torch.Tensor]: # List of (*** c) or (*** ... c) + hid_len = hid_shape.prod(-1) + hid = hid.split(hid_len.tolist()) + hid = [x.unflatten(0, s.tolist()) for x, s in zip(hid, hid_shape)] + return hid + +class NaDiT(nn.Module): + + def __init__( + self, + norm_eps, + num_layers, + mlp_type, + vid_in_channels = 33, + vid_out_channels = SEEDVR2_LATENT_CHANNELS, + vid_dim = 2560, + txt_in_dim = 5120, + heads = 20, + head_dim = 128, + mm_layers = 10, + expand_ratio = 4, + qk_bias = False, + patch_size = (1, 2, 2), + rope_dim = 128, + rope_type = "mmrope3d", + vid_out_norm: Optional[str] = None, + image_model = None, + device = None, + dtype = None, + operations = None, + ): + if image_model not in (None, "seedvr2"): + raise ValueError(f"SeedVR2 NaDiT expected image_model='seedvr2', got {image_model!r}.") + self._7b_version = vid_dim == SEEDVR2_7B_VID_DIM + if self._7b_version: + rope_type = "rope3d" + self.dtype = dtype + factory_kwargs = {"device": device, "dtype": dtype} + window_method = num_layers // 2 * ["720pwin_by_size_bysize","720pswin_by_size_bysize"] + txt_dim = vid_dim + emb_dim = vid_dim * 6 + window = num_layers * [(4,3,3)] + ada = AdaSingle + norm = operations.RMSNorm + qk_norm = operations.RMSNorm + super().__init__() + self.register_buffer("positive_conditioning", torch.empty((58, 5120), device=device, dtype=dtype)) + self.register_buffer("negative_conditioning", torch.empty((64, 5120), device=device, dtype=dtype)) + self.vid_in = NaPatchIn( + in_channels=vid_in_channels, + patch_size=patch_size, + dim=vid_dim, + device=device, dtype=dtype, operations=operations + ) + self.txt_in = ( + operations.Linear(txt_in_dim, txt_dim, **factory_kwargs) + if txt_in_dim and txt_in_dim != txt_dim + else nn.Identity() + ) + self.emb_in = TimeEmbedding( + sinusoidal_dim=BYTEDANCE_SINUSOIDAL_DIM, + hidden_dim=max(vid_dim, txt_dim), + output_dim=emb_dim, + device=device, dtype=dtype, operations=operations + ) + + if window is None or isinstance(window[0], int): + window = [window] * num_layers + + rope_dim = rope_dim if rope_dim is not None else head_dim // 2 + self.blocks = nn.ModuleList( + [ + NaMMSRTransformerBlock( + vid_dim=vid_dim, + txt_dim=txt_dim, + emb_dim=emb_dim, + heads=heads, + head_dim=head_dim, + expand_ratio=expand_ratio, + norm=norm, + norm_eps=norm_eps, + ada=ada, + qk_bias=qk_bias, + qk_norm=qk_norm, + mlp_type=mlp_type, + rope_dim = rope_dim, + window=window[i], + window_method=window_method[i], + version = self._7b_version, + is_last_layer=(i == num_layers - 1) and not self._7b_version, + rope_type = rope_type, + shared_weights=not ( + (i < mm_layers) if isinstance(mm_layers, int) else mm_layers[i] + ), + operations = operations, + **factory_kwargs + ) + for i in range(num_layers) + ] + ) + self.vid_out = NaPatchOut( + out_channels=vid_out_channels, + patch_size=patch_size, + dim=vid_dim, + device=device, dtype=dtype, operations=operations + ) + + self.vid_out_norm = None + if vid_out_norm is not None: + self.vid_out_norm = operations.RMSNorm( + normalized_shape=vid_dim, + eps=norm_eps, + elementwise_affine=True, + device=device, dtype=dtype + ) + self.vid_out_ada = ada( + dim=vid_dim, + emb_dim=emb_dim, + layers=["out"], + modes=["in"], + device=device, dtype=dtype + ) + + def _resolve_text_conditioning(self, context, cond_or_uncond=None): + if context is None or context.numel() == 0: + context = self.positive_conditioning + return flatten([context]) + if NaDiT._seedvr2_is_single_conditioning_branch(cond_or_uncond): + if context.shape[0] == 1: + context = context.squeeze(0) + return flatten([context]) + return flatten(context.unbind(0)) + if context.shape[0] % 2 != 0: + raise ValueError(f"SeedVR2 expected an even text-conditioning batch, got shape {tuple(context.shape)}") + neg_cond, pos_cond = context.chunk(2, dim=0) + if pos_cond.shape[0] == 1: + pos_cond, neg_cond = pos_cond.squeeze(0), neg_cond.squeeze(0) + return flatten([pos_cond, neg_cond]) + return flatten((*pos_cond.unbind(0), *neg_cond.unbind(0))) + + @staticmethod + def _seedvr2_is_single_conditioning_branch(cond_or_uncond): + if cond_or_uncond is None or len(cond_or_uncond) == 0: + return False + first = cond_or_uncond[0] + return all(entry == first for entry in cond_or_uncond) + + @staticmethod + def _check_seedvr2_video_latent(x, channels, name): + if x.ndim != 5: + raise ValueError(f"SeedVR2 expected {name} to be 5-D native latent, got shape {tuple(x.shape)}.") + if x.shape[1] != channels: + raise ValueError(f"SeedVR2 expected {name} channels to be {channels}, got shape {tuple(x.shape)}.") + return x + + def _swap_pos_neg_halves(self, out, cond_or_uncond=None): + if NaDiT._seedvr2_is_single_conditioning_branch(cond_or_uncond): + return out + pos, neg = out.chunk(2, dim=0) + return torch.cat([neg, pos], dim=0) + + def forward( + self, + x, + timestep, + context, # l c + disable_cache: bool = False, + **kwargs + ): + transformer_options = kwargs.get("transformer_options", {}) + patches_replace = transformer_options.get("patches_replace", {}) + blocks_replace = patches_replace.get("dit", {}) + conditions = kwargs.get("condition") + if conditions is None: + raise ValueError("SeedVR2 requires conditioning latents from the SeedVR2Conditioning node.") + x = self._check_seedvr2_video_latent(x, SEEDVR2_LATENT_CHANNELS, "latent") + conditions = self._check_seedvr2_video_latent(conditions, SEEDVR2_LATENT_CHANNELS + 1, "conditioning") + b, _, t, h, w = x.shape + if conditions.shape[0] != b or conditions.shape[2:] != (t, h, w): + raise ValueError( + f"SeedVR2 conditioning shape must match latent batch/temporal/spatial dimensions; got latent {tuple(x.shape)} and conditioning {tuple(conditions.shape)}." + ) + x = x.movedim(1, -1) + conditions = conditions.movedim(1, -1) + cache = Cache(disable=disable_cache) + + txt, txt_shape = self._resolve_text_conditioning(context, transformer_options.get("cond_or_uncond")) + + vid, vid_shape = flatten(x) + cond_latent, _ = flatten(conditions) + + vid = torch.cat([vid, cond_latent], dim=-1) + + txt = self.txt_in(txt) + + vid_shape_before_patchify = vid_shape + vid, vid_shape = self.vid_in(vid, vid_shape, cache=cache) + + emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype) + + for i, block in enumerate(self.blocks): + if ("block", i) in blocks_replace: + def block_wrap(args): + out = {} + out["vid"], out["txt"], out["vid_shape"], out["txt_shape"] = block( + vid=args["vid"], + txt=args["txt"], + vid_shape=args["vid_shape"], + txt_shape=args["txt_shape"], + emb=args["emb"], + cache=args["cache"], + ) + return out + out = blocks_replace[("block", i)]({ + "vid":vid, + "txt":txt, + "vid_shape":vid_shape, + "txt_shape":txt_shape, + "emb":emb, + "cache":cache, + }, {"original_block": block_wrap}) + vid, txt, vid_shape, txt_shape = out["vid"], out["txt"], out["vid_shape"], out["txt_shape"] + else: + vid, txt, vid_shape, txt_shape = block( + vid=vid, + txt=txt, + vid_shape=vid_shape, + txt_shape=txt_shape, + emb=emb, + cache=cache, + ) + + if self.vid_out_norm: + vid = self.vid_out_norm(vid) + vid = self.vid_out_ada( + vid, + emb=emb, + layer="out", + mode="in", + hid_len=cache("vid_len", lambda: vid_shape.prod(-1)), + cache=cache, + branch_tag="vid", + ) + + vid, vid_shape = self.vid_out(vid, vid_shape, cache, vid_shape_before_patchify = vid_shape_before_patchify) + vid = unflatten(vid, vid_shape) + out = torch.stack(vid) + out = out.movedim(-1, 1) + return self._swap_pos_neg_halves(out, transformer_options.get("cond_or_uncond")) diff --git a/comfy/ldm/seedvr/vae.py b/comfy/ldm/seedvr/vae.py new file mode 100644 index 000000000..c9f430184 --- /dev/null +++ b/comfy/ldm/seedvr/vae.py @@ -0,0 +1,1612 @@ +from typing import Literal, Optional, Tuple +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from contextlib import contextmanager +from comfy.utils import ProgressBar + +from comfy.ldm.seedvr.constants import ( + BYTEDANCE_BLOCK_OUT_CHANNELS, + BYTEDANCE_GN_CHUNKS_FP16, + BYTEDANCE_GN_CHUNKS_FP32, + BYTEDANCE_LOGVAR_CLAMP_MAX, + BYTEDANCE_LOGVAR_CLAMP_MIN, + BYTEDANCE_SLICING_SAMPLE_MIN, + BYTEDANCE_VAE_CONV_MEM_GIB, + BYTEDANCE_VAE_NORM_MEM_GIB, + BYTEDANCE_VAE_SCALING_FACTOR, + BYTEDANCE_VAE_SHIFTING_FACTOR, + BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE, + BYTEDANCE_VAE_TEMPORAL_DOWNSAMPLE, + SEEDVR2_LATENT_CHANNELS, +) +from comfy.ldm.modules.attention import optimized_attention +from comfy.ldm.modules.diffusionmodules.model import vae_attention + +import math +from enum import Enum + +import logging +import comfy.model_management +import comfy.ops +ops = comfy.ops.disable_weight_init + + +def _seedvr2_temporal_slicing_min_size(temporal_size, temporal_overlap, temporal_scale=1): + if temporal_size is None: + return None + + temporal_size = int(temporal_size) + if temporal_size <= 0: + return None + + temporal_overlap = max(0, int(temporal_overlap or 0)) + temporal_overlap = min(temporal_overlap, temporal_size - 1) + temporal_step = temporal_size - temporal_overlap + temporal_scale = max(1, int(temporal_scale)) + return max(1, math.ceil(temporal_step / temporal_scale)) + + +def _seedvr2_clamped_spatial_overlap(overlap, tile_size): + overlap = max(0, int(overlap)) + tile_size = max(1, int(tile_size)) + return min(overlap, tile_size - 1) + + +def tiled_vae( + x, + vae_model, + tile_size=(512, 512), + tile_overlap=(64, 64), + temporal_size=16, + temporal_overlap=0, + encode=True, +): + if x.ndim != 5: + x = x.unsqueeze(2) + + _, _, d, h, w = x.shape + + sf_s = getattr(vae_model, "spatial_downsample_factor", BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE) + sf_t = getattr(vae_model, "temporal_downsample_factor", BYTEDANCE_VAE_TEMPORAL_DOWNSAMPLE) + if encode: + slicing_attr = "slicing_sample_min_size" + slicing_min_size = _seedvr2_temporal_slicing_min_size(temporal_size, temporal_overlap) + else: + slicing_attr = "slicing_latent_min_size" + slicing_min_size = _seedvr2_temporal_slicing_min_size(temporal_size, temporal_overlap, sf_t) + if encode: + ti_h, ti_w = tile_size + ov_h = _seedvr2_clamped_spatial_overlap(tile_overlap[0], ti_h) + ov_w = _seedvr2_clamped_spatial_overlap(tile_overlap[1], ti_w) + blend_ov_h = max(0, ov_h // sf_s) + blend_ov_w = max(0, ov_w // sf_s) + target_d = (d + sf_t - 1) // sf_t + target_h = (h + sf_s - 1) // sf_s + target_w = (w + sf_s - 1) // sf_s + else: + ti_h = max(1, tile_size[0] // sf_s) + ti_w = max(1, tile_size[1] // sf_s) + ov_h = _seedvr2_clamped_spatial_overlap(tile_overlap[0] // sf_s, ti_h) + ov_w = _seedvr2_clamped_spatial_overlap(tile_overlap[1] // sf_s, ti_w) + blend_ov_h = ov_h * sf_s + blend_ov_w = ov_w * sf_s + + target_d = max(1, d * sf_t - (sf_t - 1)) + target_h = h * sf_s + target_w = w * sf_s + + stride_h = max(1, ti_h - ov_h) + stride_w = max(1, ti_w - ov_w) + + storage_device = vae_model.device + result = None + count = None + def run_temporal_chunks(spatial_tile, model=vae_model, device=storage_device): + device = torch.device(device) + t_chunk = spatial_tile.to(device=device, dtype=next(model.parameters()).dtype, non_blocking=True).contiguous() + old_device = getattr(model, "device", None) + model.device = device + old_slicing_min_size = getattr(model, slicing_attr, None) + if old_slicing_min_size is not None and slicing_min_size is not None: + if slicing_min_size <= 0: + setattr(model, slicing_attr, t_chunk.shape[2]) + else: + setattr(model, slicing_attr, slicing_min_size) + try: + if encode: + out = model.encode(t_chunk) + else: + out = model.decode_(t_chunk) + finally: + if old_slicing_min_size is not None and slicing_min_size is not None: + setattr(model, slicing_attr, old_slicing_min_size) + if old_device is not None: + model.device = old_device + if out.ndim == 4: + out = out.unsqueeze(2) + return out.to(storage_device) + + ramp_cache = {} + def get_ramp(steps): + if steps not in ramp_cache: + t = torch.linspace(0, 1, steps=steps, device=storage_device, dtype=torch.float32) + ramp_cache[steps] = 0.5 - 0.5 * torch.cos(t * torch.pi) + return ramp_cache[steps] + + tile_ranges = [] + for y_idx in range(0, h, stride_h): + y_end = min(y_idx + ti_h, h) + if y_idx > 0 and (y_end - y_idx) <= ov_h: + continue + for x_idx in range(0, w, stride_w): + x_end = min(x_idx + ti_w, w) + if x_idx > 0 and (x_end - x_idx) <= ov_w: + continue + tile_ranges.append((y_idx, y_end, x_idx, x_end)) + + total_tiles = len(tile_ranges) + bar = ProgressBar(total_tiles) + single_spatial_tile = h <= ti_h and w <= ti_w + + def run_tile(tile_index, tile_range): + y_idx, y_end, x_idx, x_end = tile_range + tile_x = x[:, :, :, y_idx:y_end, x_idx:x_end] + tile_out = run_temporal_chunks(tile_x) + return tile_index, y_idx, y_end, x_idx, x_end, tile_out + + ordered_tile_outputs = ( + run_tile(tile_index, tile_range) + for tile_index, tile_range in enumerate(tile_ranges) + ) + + for _, y_idx, y_end, x_idx, x_end, tile_out in ordered_tile_outputs: + + if single_spatial_tile: + result = tile_out[:, :, :target_d, :target_h, :target_w] + if result.device != x.device or result.dtype != x.dtype: + result = result.to(device=x.device, dtype=x.dtype) + if x.shape[2] == 1 and sf_t == 1: + result = result.squeeze(2) + bar.update(1) + return result + + if result is None: + b_out, c_out = tile_out.shape[0], tile_out.shape[1] + result = torch.zeros((b_out, c_out, target_d, target_h, target_w), device=storage_device, dtype=torch.float32) + count = torch.zeros((1, 1, 1, target_h, target_w), device=storage_device, dtype=torch.float32) + + if encode: + ys, ye = y_idx // sf_s, (y_idx // sf_s) + tile_out.shape[3] + xs, xe = x_idx // sf_s, (x_idx // sf_s) + tile_out.shape[4] + cur_ov_h = max(0, min(blend_ov_h, tile_out.shape[3] // 2)) + cur_ov_w = max(0, min(blend_ov_w, tile_out.shape[4] // 2)) + else: + ys, ye = y_idx * sf_s, (y_idx * sf_s) + tile_out.shape[3] + xs, xe = x_idx * sf_s, (x_idx * sf_s) + tile_out.shape[4] + cur_ov_h = max(0, min(blend_ov_h, tile_out.shape[3] // 2)) + cur_ov_w = max(0, min(blend_ov_w, tile_out.shape[4] // 2)) + + w_h = torch.ones((tile_out.shape[3],), device=storage_device) + w_w = torch.ones((tile_out.shape[4],), device=storage_device) + + if cur_ov_h > 0: + r = get_ramp(cur_ov_h) + if y_idx > 0: + w_h[:cur_ov_h] = r + if y_end < h: + w_h[-cur_ov_h:] = 1.0 - r + + if cur_ov_w > 0: + r = get_ramp(cur_ov_w) + if x_idx > 0: + w_w[:cur_ov_w] = r + if x_end < w: + w_w[-cur_ov_w:] = 1.0 - r + + final_weight = w_h.view(1,1,1,-1,1) * w_w.view(1,1,1,1,-1) + + valid_d = min(tile_out.shape[2], result.shape[2]) + tile_out = tile_out[:, :, :valid_d, :, :] + + tile_out.mul_(final_weight) + + result[:, :, :valid_d, ys:ye, xs:xe] += tile_out + count[:, :, :, ys:ye, xs:xe] += final_weight + + del tile_out, final_weight, w_h, w_w + bar.update(1) + + result.div_(count.clamp(min=1e-6)) + + if result.device != x.device or result.dtype != x.dtype: + result = result.to(device=x.device, dtype=x.dtype) + + if x.shape[2] == 1 and sf_t == 1: + result = result.squeeze(2) + + return result + +_NORM_LIMIT = float("inf") +def get_norm_limit(): + return _NORM_LIMIT + + +def set_norm_limit(value: Optional[float] = None): + global _NORM_LIMIT + if value is None: + value = float("inf") + _NORM_LIMIT = value + +@contextmanager +def ignore_padding(model): + orig_padding = model.padding + model.padding = (0, 0, 0) + try: + yield + finally: + model.padding = orig_padding + +class MemoryState(Enum): + DISABLED = 0 + INITIALIZING = 1 + ACTIVE = 2 + UNSET = 3 + +def get_cache_size(conv_module, input_len, pad_len, dim=0): + dilated_kernel_size = conv_module.dilation[dim] * (conv_module.kernel_size[dim] - 1) + 1 + output_len = (input_len + pad_len - dilated_kernel_size) // conv_module.stride[dim] + 1 + remain_len = ( + input_len + pad_len - ((output_len - 1) * conv_module.stride[dim] + dilated_kernel_size) + ) + overlap_len = dilated_kernel_size - conv_module.stride[dim] + cache_len = overlap_len + remain_len + + if output_len <= 0: + raise ValueError( + f"SeedVR2 VAE cache input is too short for convolution: input_len={input_len}, pad_len={pad_len}." + ) + return cache_len + +class DiagonalGaussianDistribution(object): + def __init__(self, parameters: torch.Tensor): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, BYTEDANCE_LOGVAR_CLAMP_MIN, BYTEDANCE_LOGVAR_CLAMP_MAX) + + def mode(self): + return self.mean + +class SpatialNorm(nn.Module): + def __init__( + self, + f_channels: int, + zq_channels: int, + ): + super().__init__() + self.norm_layer = ops.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) + self.conv_y = ops.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) + self.conv_b = ops.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor: + f_size = f.shape[-2:] + zq = F.interpolate(zq, size=f_size, mode="nearest") + norm_f = self.norm_layer(f) + new_f = norm_f * self.conv_y(zq) + self.conv_b(zq) + return new_f + +class Attention(nn.Module): + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + bias: bool = False, + norm_num_groups: Optional[int] = None, + spatial_norm_dim: Optional[int] = None, + out_bias: bool = True, + eps: float = 1e-5, + rescale_output_factor: float = 1.0, + residual_connection: bool = False, + ): + super().__init__() + + self.inner_dim = dim_head * heads + self.rescale_output_factor = rescale_output_factor + self.residual_connection = residual_connection + self.out_dim = query_dim + self.heads = heads + + if norm_num_groups is not None: + self.group_norm = ops.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True) + else: + self.group_norm = None + + if spatial_norm_dim is not None: + self.spatial_norm = SpatialNorm(f_channels=query_dim, zq_channels=spatial_norm_dim) + else: + self.spatial_norm = None + + self.to_q = ops.Linear(query_dim, self.inner_dim, bias=bias) + self.to_k = ops.Linear(query_dim, self.inner_dim, bias=bias) + self.to_v = ops.Linear(query_dim, self.inner_dim, bias=bias) + self.to_out = nn.ModuleList([]) + self.to_out.append(ops.Linear(self.inner_dim, self.out_dim, bias=out_bias)) + self.to_out.append(nn.Identity()) + + self.optimized_vae_attention = vae_attention() + + def forward( + self, + hidden_states: torch.Tensor, + temb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + residual = hidden_states + if self.spatial_norm is not None: + hidden_states = self.spatial_norm(hidden_states, temb) + + input_ndim = hidden_states.ndim + + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + batch_size = hidden_states.shape[0] + + if self.group_norm is not None: + hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + inner_dim = key.shape[-1] + head_dim = inner_dim // self.heads + + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + if input_ndim == 4 and self.heads == 1: + query = query.squeeze(1).transpose(1, 2).reshape(batch_size, head_dim, height, width) + key = key.squeeze(1).transpose(1, 2).reshape(batch_size, head_dim, height, width) + value = value.squeeze(1).transpose(1, 2).reshape(batch_size, head_dim, height, width) + hidden_states = self.optimized_vae_attention(query, key, value).reshape(batch_size, self.heads, head_dim, height * width).transpose(2, 3) + else: + hidden_states = optimized_attention(query, key, value, heads = self.heads, skip_reshape=True, skip_output_reshape=True) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + + if input_ndim == 4: + hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + + if self.residual_connection: + hidden_states = hidden_states + residual + + hidden_states = hidden_states / self.rescale_output_factor + + return hidden_states + + +def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor: + input_dtype = x.dtype + if isinstance(norm_layer, (ops.LayerNorm, ops.RMSNorm)): + if x.ndim == 4: + x = x.permute(0, 2, 3, 1) + x = norm_layer(x) + x = x.permute(0, 3, 1, 2) + return x.to(input_dtype) + if x.ndim == 5: + x = x.permute(0, 2, 3, 4, 1) + x = norm_layer(x) + x = x.permute(0, 4, 1, 2, 3) + return x.to(input_dtype) + if isinstance(norm_layer, (ops.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)): + if x.ndim <= 4: + return norm_layer(x).to(input_dtype) + if x.ndim == 5: + b, c, t, h, w = x.shape + x = x.transpose(1, 2).reshape(b * t, c, h, w) + memory_occupy = x.numel() * x.element_size() / 1024**3 + if isinstance(norm_layer, ops.GroupNorm) and memory_occupy > get_norm_limit(): + num_chunks = min(BYTEDANCE_GN_CHUNKS_FP16 if x.element_size() == 2 else BYTEDANCE_GN_CHUNKS_FP32, norm_layer.num_groups) + if norm_layer.num_groups % num_chunks != 0: + raise ValueError( + f"SeedVR2 VAE GroupNorm groups must divide chunks: groups={norm_layer.num_groups}, chunks={num_chunks}." + ) + num_groups_per_chunk = norm_layer.num_groups // num_chunks + + x = list(x.chunk(num_chunks, dim=1)) + weights = norm_layer.weight.chunk(num_chunks, dim=0) + biases = norm_layer.bias.chunk(num_chunks, dim=0) + for i, (w, bias) in enumerate(zip(weights, biases)): + x[i] = F.group_norm(x[i], num_groups_per_chunk, w, bias, norm_layer.eps) + x[i] = x[i].to(input_dtype) + x = torch.cat(x, dim=1) + else: + x = norm_layer(x) + x = x.reshape((b, t, x.size(1), x.size(2), x.size(3))).transpose(1, 2) + return x.to(input_dtype) + raise TypeError(f"SeedVR2 VAE unsupported norm layer type: {type(norm_layer).__name__}") + +_receptive_field_t = Literal["half", "full"] + +def extend_head(tensor, times: int = 2, memory = None): + if memory is not None: + return torch.cat((memory.to(tensor), tensor), dim=2) + if times < 0: + raise ValueError(f"SeedVR2 VAE extend_head expected times >= 0, got {times}.") + if times == 0: + return tensor + else: + tile_repeat = [1] * tensor.ndim + tile_repeat[2] = times + return torch.cat(tensors=(torch.tile(tensor[:, :, :1], tile_repeat), tensor), dim=2) + +def cache_send_recv(tensor, cache_size, times, memory=None): + recv_buffer = None + + if memory is not None: + recv_buffer = memory.to(tensor[0]) + elif times > 0: + tile_repeat = [1] * tensor[0].ndim + tile_repeat[2] = times + recv_buffer = torch.tile(tensor[0][:, :, :1], tile_repeat) + + return recv_buffer + +class InflatedCausalConv3d(ops.Conv3d): + def __init__( + self, + *args, + inflation_mode, + **kwargs, + ): + self.inflation_mode = inflation_mode + super().__init__(*args, **kwargs) + self.temporal_padding = self.padding[0] + self.padding = (0, *self.padding[1:]) + self.memory_limit = float("inf") + self.logged_once = False + + def set_memory_limit(self, value: float): + self.memory_limit = value + + def _conv_forward(self, input, weight, bias, *args, **kwargs): + try: + return super()._conv_forward(input, weight, bias, *args, **kwargs) + except NotImplementedError: + # for: Could not run 'aten::cudnn_convolution' with arguments from the 'CPU' backend + if not self.logged_once: + logging.warning("VAE is on CPU for decoding. This is most likely due to not enough memory") + self.logged_once = True + return F.conv3d(input, weight, bias, *args, **kwargs) + + def memory_limit_conv( + self, + x, + *, + split_dim=3, + padding=(0, 0, 0, 0, 0, 0), + prev_cache=None, + ): + if math.isinf(self.memory_limit): + if prev_cache is not None: + x = torch.cat([prev_cache, x], dim=split_dim - 1) + return super().forward(x) + + shape = list(x.size()) + if prev_cache is not None: + shape[split_dim - 1] += prev_cache.size(split_dim - 1) + for i, pad_sum in enumerate((padding[4] + padding[5], padding[2] + padding[3], padding[0] + padding[1])): + shape[-3 + i] += pad_sum + memory_occupy = math.prod(shape) * x.element_size() / 1024**3 # GiB + if memory_occupy < self.memory_limit or split_dim == x.ndim: + x_concat = x + if prev_cache is not None: + x_concat = torch.cat([prev_cache, x], dim=split_dim - 1) + + def pad_and_forward(): + padded = F.pad(x_concat, padding, mode='constant', value=0.0) + if not padded.is_contiguous(): + padded = padded.contiguous() + with ignore_padding(self): + return torch.nn.Conv3d.forward(self, padded) + + return pad_and_forward() + + num_splits = math.ceil(memory_occupy / self.memory_limit) + size_per_split = x.size(split_dim) // num_splits + split_sizes = [size_per_split] * (num_splits - 1) + split_sizes += [x.size(split_dim) - sum(split_sizes)] + + x = list(x.split(split_sizes, dim=split_dim)) + if prev_cache is not None: + prev_cache = list(prev_cache.split(split_sizes, dim=split_dim)) + cache = None + for idx in range(len(x)): + if prev_cache is not None: + x[idx] = torch.cat([prev_cache[idx], x[idx]], dim=split_dim - 1) + + lpad_dim = (x[idx].ndim - split_dim - 1) * 2 + rpad_dim = lpad_dim + 1 + padding = list(padding) + padding[lpad_dim] = self.padding[split_dim - 2] if idx == 0 else 0 + padding[rpad_dim] = self.padding[split_dim - 2] if idx == len(x) - 1 else 0 + pad_len = padding[lpad_dim] + padding[rpad_dim] + padding = tuple(padding) + + next_cache = None + cache_len = cache.size(split_dim) if cache is not None else 0 + next_cache_size = get_cache_size( + conv_module=self, + input_len=x[idx].size(split_dim) + cache_len, + pad_len=pad_len, + dim=split_dim - 2, + ) + if next_cache_size != 0: + if next_cache_size > x[idx].size(split_dim): + raise ValueError( + f"SeedVR2 VAE cache size {next_cache_size} exceeds split size {x[idx].size(split_dim)}." + ) + next_cache = ( + x[idx].transpose(0, split_dim)[-next_cache_size:].transpose(0, split_dim) + ) + + x[idx] = self.memory_limit_conv( + x[idx], + split_dim=split_dim + 1, + padding=padding, + prev_cache=cache + ) + + cache = next_cache + + output = torch.cat(x, dim=split_dim) + return output + + def forward( + self, + input, + memory_state: MemoryState = MemoryState.UNSET, + memory_cache = None, + ) -> Tensor: + if memory_state == MemoryState.UNSET: + raise ValueError("SeedVR2 VAE convolution requires an explicit MemoryState.") + if memory_cache is None: + memory_cache = {} + if memory_state != MemoryState.ACTIVE: + memory_cache.pop(self, None) + if ( + math.isinf(self.memory_limit) + and torch.is_tensor(input) + ): + return self.basic_forward(input, memory_state, memory_cache) + return self.slicing_forward(input, memory_state, memory_cache) + + def basic_forward(self, input: Tensor, memory_state: MemoryState = MemoryState.UNSET, memory_cache = None): + mem_size = self.stride[0] - self.kernel_size[0] + memory = memory_cache.get(self) if memory_cache is not None else None + if (memory is not None) and (memory_state == MemoryState.ACTIVE): + input = extend_head(input, memory=memory, times=-1) + else: + input = extend_head(input, times=self.temporal_padding * 2) + next_memory = ( + input[:, :, mem_size:].detach() + if (mem_size != 0 and memory_state != MemoryState.DISABLED) + else None + ) + if memory_cache is not None and memory_state != MemoryState.DISABLED: + if next_memory is None: + memory_cache.pop(self, None) + else: + memory_cache[self] = next_memory + return super().forward(input) + + def slicing_forward( + self, + input, + memory_state: MemoryState = MemoryState.UNSET, + memory_cache = None, + ) -> Tensor: + if memory_cache is None: + memory_cache = {} + squeeze_out = False + if torch.is_tensor(input): + input = [input] + squeeze_out = True + + cache_size = self.kernel_size[0] - self.stride[0] + memory = memory_cache.get(self) if memory_cache is not None else None + cache = cache_send_recv( + input, cache_size=cache_size, memory=memory, times=self.temporal_padding * 2 + ) + + if ( + memory_state in [MemoryState.INITIALIZING, MemoryState.ACTIVE] + and cache_size != 0 + ): + if cache_size > input[-1].size(2) and cache is not None and len(input) == 1: + input[0] = torch.cat([cache, input[0]], dim=2) + cache = None + if cache_size <= input[-1].size(2): + memory_cache[self] = input[-1][:, :, -cache_size:].detach().contiguous() + + padding = tuple(x for x in reversed(self.padding) for _ in range(2)) + for i in range(len(input)): + next_cache = None + cache_size = 0 + if i < len(input) - 1: + cache_len = cache.size(2) if cache is not None else 0 + cache_size = get_cache_size(self, input[i].size(2) + cache_len, pad_len=0) + if cache_size != 0: + if cache_size > input[i].size(2) and cache is not None: + input[i] = torch.cat([cache, input[i]], dim=2) + cache = None + if cache_size > input[i].size(2): + raise ValueError(f"SeedVR2 VAE cache size {cache_size} exceeds input length {input[i].size(2)}.") + next_cache = input[i][:, :, -cache_size:] + + input[i] = self.memory_limit_conv( + input[i], + padding=padding, + prev_cache=cache + ) + + cache = next_cache + + return input[0] if squeeze_out else input + +def remove_head(tensor: Tensor, times: int = 1) -> Tensor: + if times == 0: + return tensor + return torch.cat(tensors=(tensor[:, :, :1], tensor[:, :, times + 1 :]), dim=2) + +class Upsample3D(nn.Module): + + def __init__( + self, + channels, + out_channels = None, + inflation_mode = "tail", + temporal_up: bool = False, + spatial_up: bool = True, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + + conv = InflatedCausalConv3d( + self.channels, + self.out_channels, + 3, + padding=1, + inflation_mode=inflation_mode, + ) + + self.temporal_up = temporal_up + self.spatial_up = spatial_up + self.temporal_ratio = 2 if temporal_up else 1 + self.spatial_ratio = 2 if spatial_up else 1 + + upscale_ratio = (self.spatial_ratio**2) * self.temporal_ratio + self.upscale_conv = ops.Conv3d( + self.channels, self.channels * upscale_ratio, kernel_size=1, padding=0 + ) + + self.conv = conv + + def forward( + self, + hidden_states: torch.FloatTensor, + memory_state=None, + memory_cache=None, + ) -> torch.FloatTensor: + if hidden_states.shape[1] != self.channels: + raise ValueError(f"SeedVR2 upsample expected {self.channels} channels, got {hidden_states.shape[1]}.") + + hidden_states = self.upscale_conv(hidden_states) + b, channels, f, h, w = hidden_states.shape + c = channels // (self.spatial_ratio * self.spatial_ratio * self.temporal_ratio) + hidden_states = hidden_states.view(b, self.spatial_ratio, self.spatial_ratio, self.temporal_ratio, c, f, h, w) + hidden_states = hidden_states.permute(0, 4, 5, 3, 6, 1, 7, 2).reshape( + b, + c, + f * self.temporal_ratio, + h * self.spatial_ratio, + w * self.spatial_ratio, + ) + + if self.temporal_up and memory_state != MemoryState.ACTIVE: + hidden_states = remove_head(hidden_states) + + hidden_states = self.conv(hidden_states, memory_state=memory_state, memory_cache=memory_cache) + + return hidden_states + + +class Downsample3D(nn.Module): + def __init__( + self, + channels, + out_channels = None, + inflation_mode = "tail", + spatial_down: bool = False, + temporal_down: bool = False, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.temporal_down = temporal_down + self.spatial_down = spatial_down + + self.temporal_ratio = 2 if temporal_down else 1 + self.spatial_ratio = 2 if spatial_down else 1 + + self.temporal_kernel = 3 if temporal_down else 1 + self.spatial_kernel = 3 if spatial_down else 1 + + self.conv = InflatedCausalConv3d( + self.channels, + self.out_channels, + kernel_size=(self.temporal_kernel, self.spatial_kernel, self.spatial_kernel), + stride=(self.temporal_ratio, self.spatial_ratio, self.spatial_ratio), + padding=(1 if self.temporal_down else 0, 0, 0), + inflation_mode=inflation_mode, + ) + + + def forward( + self, + hidden_states: torch.FloatTensor, + memory_state = None, + memory_cache = None, + ) -> torch.FloatTensor: + + if hidden_states.shape[1] != self.channels: + raise ValueError(f"SeedVR2 downsample expected {self.channels} channels, got {hidden_states.shape[1]}.") + + if self.spatial_down: + pad = (0, 1, 0, 1) + hidden_states = F.pad(hidden_states, pad, mode="constant", value=0) + + if hidden_states.shape[1] != self.channels: + raise ValueError(f"SeedVR2 downsample expected {self.channels} channels after padding, got {hidden_states.shape[1]}.") + + hidden_states = self.conv(hidden_states, memory_state=memory_state, memory_cache=memory_cache) + + return hidden_states + + +class ResnetBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + temb_channels: int = 512, + groups: int = 32, + groups_out: Optional[int] = None, + eps: float = 1e-6, + output_scale_factor: float = 1.0, + skip_time_act: bool = False, + inflation_mode = "tail", + time_receptive_field: _receptive_field_t = "half", + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.output_scale_factor = output_scale_factor + self.skip_time_act = skip_time_act + self.nonlinearity = nn.SiLU() + if temb_channels is not None: + self.time_emb_proj = ops.Linear(temb_channels, self.out_channels) + else: + self.time_emb_proj = None + self.norm1 = ops.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + if groups_out is None: + groups_out = groups + self.norm2 = ops.GroupNorm(num_groups=groups_out, num_channels=self.out_channels, eps=eps, affine=True) + self.use_in_shortcut = self.in_channels != self.out_channels + self.conv1 = InflatedCausalConv3d( + self.in_channels, + self.out_channels, + kernel_size=(1, 3, 3) if time_receptive_field == "half" else (3, 3, 3), + stride=1, + padding=(0, 1, 1) if time_receptive_field == "half" else (1, 1, 1), + inflation_mode=inflation_mode, + ) + + self.conv2 = InflatedCausalConv3d( + self.out_channels, + self.out_channels, + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = InflatedCausalConv3d( + self.in_channels, + self.out_channels, + kernel_size=1, + stride=1, + padding=0, + bias=True, + inflation_mode=inflation_mode, + ) + + def forward(self, input_tensor, temb, memory_state = None, memory_cache = None): + hidden_states = input_tensor + + hidden_states = causal_norm_wrapper(self.norm1, hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv1(hidden_states, memory_state=memory_state, memory_cache=memory_cache) + + if self.time_emb_proj is not None: + if not self.skip_time_act: + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb)[:, :, None, None] + + if temb is not None: + hidden_states = hidden_states + temb + + hidden_states = causal_norm_wrapper(self.norm2, hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv2(hidden_states, memory_state=memory_state, memory_cache=memory_cache) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor, memory_state=memory_state, memory_cache=memory_cache) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + return output_tensor + + +class DownEncoderBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_groups: int = 32, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + inflation_mode = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_down: bool = True, + spatial_down: bool = True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock3D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + output_scale_factor=output_scale_factor, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample3D( + out_channels, + out_channels=out_channels, + temporal_down=temporal_down, + spatial_down=spatial_down, + inflation_mode=inflation_mode, + ) + ] + ) + else: + self.downsamplers = None + + def forward( + self, + hidden_states: torch.FloatTensor, + memory_state = None, + memory_cache = None, + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None, memory_state=memory_state, memory_cache=memory_cache) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, memory_state=memory_state, memory_cache=memory_cache) + + return hidden_states + + +class UpDecoderBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_groups: int = 32, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + temb_channels: Optional[int] = None, + inflation_mode = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_up: bool = True, + spatial_up: bool = True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlock3D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + output_scale_factor=output_scale_factor, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + Upsample3D( + out_channels, + out_channels=out_channels, + temporal_up=temporal_up, + spatial_up=spatial_up, + inflation_mode=inflation_mode, + ) + ] + ) + else: + self.upsamplers = None + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + memory_state=None, + memory_cache=None, + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None, memory_state=memory_state, memory_cache=memory_cache) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, memory_state=memory_state, memory_cache=memory_cache) + + return hidden_states + + +class UNetMidBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_groups: int = 32, + add_attention: bool = True, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + inflation_mode = "tail", + time_receptive_field: _receptive_field_t = "half", + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + + resnets = [ + ResnetBlock3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + output_scale_factor=output_scale_factor, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ] + attentions = [] + + if attention_head_dim is None: + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=( + resnet_groups if resnet_time_scale_shift == "default" else None + ), + spatial_norm_dim=( + temb_channels if resnet_time_scale_shift == "spatial" else None + ), + residual_connection=True, + bias=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + ResnetBlock3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + output_scale_factor=output_scale_factor, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states, temb=None, memory_state=None, memory_cache=None): + video_length = hidden_states.size(2) + hidden_states = self.resnets[0](hidden_states, temb, memory_state=memory_state, memory_cache=memory_cache) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + b, c, f, h, w = hidden_states.shape + hidden_states = hidden_states.transpose(1, 2).reshape(b * f, c, h, w) + hidden_states = attn(hidden_states, temb=temb) + hidden_states = hidden_states.reshape(b, video_length, c, h, w).transpose(1, 2) + hidden_states = resnet(hidden_states, temb, memory_state=memory_state, memory_cache=memory_cache) + + return hidden_states + + +class Encoder3D(nn.Module): + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ("DownEncoderBlock3D",), + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + norm_num_groups: int = 32, + mid_block_add_attention=True, + temporal_down_num: int = 2, + inflation_mode = "tail", + time_receptive_field: _receptive_field_t = "half", + ): + super().__init__() + self.layers_per_block = layers_per_block + self.temporal_down_num = temporal_down_num + + self.conv_in = InflatedCausalConv3d( + in_channels, + block_out_channels[0], + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + is_temporal_down_block = i >= len(block_out_channels) - self.temporal_down_num - 1 + + if down_block_type != "DownEncoderBlock3D": + raise ValueError(f"SeedVR2 encoder only supports DownEncoderBlock3D, got {down_block_type}.") + + down_block = DownEncoderBlock3D( + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=not is_final_block, + resnet_eps=1e-6, + resnet_groups=norm_num_groups, + temporal_down=is_temporal_down_block, + spatial_down=True, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + self.down_blocks.append(down_block) + + self.mid_block = UNetMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + self.conv_norm_out = ops.GroupNorm( + num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6 + ) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels + self.conv_out = InflatedCausalConv3d( + block_out_channels[-1], conv_out_channels, 3, padding=1, inflation_mode=inflation_mode + ) + + + def forward( + self, + sample: torch.FloatTensor, + memory_state = None, + memory_cache = None, + ) -> torch.FloatTensor: + sample = sample.to(next(self.parameters()).device) + sample = self.conv_in(sample, memory_state=memory_state, memory_cache=memory_cache) + for down_block in self.down_blocks: + sample = down_block(sample, memory_state=memory_state, memory_cache=memory_cache) + + sample = self.mid_block(sample, memory_state=memory_state, memory_cache=memory_cache) + + sample = causal_norm_wrapper(self.conv_norm_out, sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample, memory_state=memory_state, memory_cache=memory_cache) + + return sample + + +class Decoder3D(nn.Module): + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ("UpDecoderBlock3D",), + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + norm_num_groups: int = 32, + mid_block_add_attention=True, + inflation_mode = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_up_num: int = 2, + ): + super().__init__() + self.layers_per_block = layers_per_block + self.temporal_up_num = temporal_up_num + + self.conv_in = InflatedCausalConv3d( + in_channels, + block_out_channels[-1], + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + temb_channels = None + + self.mid_block = UNetMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=temb_channels, + add_attention=mid_block_add_attention, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + + is_final_block = i == len(block_out_channels) - 1 + is_temporal_up_block = i < self.temporal_up_num + if up_block_type != "UpDecoderBlock3D": + raise ValueError(f"SeedVR2 decoder only supports UpDecoderBlock3D, got {up_block_type}.") + up_block = UpDecoderBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + add_upsample=not is_final_block, + resnet_eps=1e-6, + resnet_groups=norm_num_groups, + temb_channels=temb_channels, + temporal_up=is_temporal_up_block, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + self.conv_norm_out = ops.GroupNorm( + num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6 + ) + self.conv_act = nn.SiLU() + self.conv_out = InflatedCausalConv3d( + block_out_channels[0], out_channels, 3, padding=1, inflation_mode=inflation_mode + ) + + + def forward( + self, + sample: torch.FloatTensor, + latent_embeds: Optional[torch.FloatTensor] = None, + memory_state = None, + memory_cache = None, + ) -> torch.FloatTensor: + + sample = sample.to(next(self.parameters()).device) + sample = self.conv_in(sample, memory_state=memory_state, memory_cache=memory_cache) + + upscale_dtype = next(iter(self.up_blocks.parameters())).dtype + sample = self.mid_block(sample, latent_embeds, memory_state=memory_state, memory_cache=memory_cache) + sample = sample.to(upscale_dtype) + + for up_block in self.up_blocks: + sample = up_block(sample, latent_embeds, memory_state=memory_state, memory_cache=memory_cache) + + sample = causal_norm_wrapper(self.conv_norm_out, sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample, memory_state=memory_state, memory_cache=memory_cache) + + return sample + +class VideoAutoencoderKL(nn.Module): + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + layers_per_block: int = 2, + latent_channels: int = SEEDVR2_LATENT_CHANNELS, + norm_num_groups: int = 32, + temporal_scale_num: int = 2, + inflation_mode = "pad", + time_receptive_field: _receptive_field_t = "full", + slicing_sample_min_size = BYTEDANCE_SLICING_SAMPLE_MIN, + ): + self.slicing_sample_min_size = slicing_sample_min_size + self.slicing_latent_min_size = slicing_sample_min_size // (2**temporal_scale_num) + block_out_channels = BYTEDANCE_BLOCK_OUT_CHANNELS + down_block_types = ("DownEncoderBlock3D",) * 4 + up_block_types = ("UpDecoderBlock3D",) * 4 + super().__init__() + + self.encoder = Encoder3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + temporal_down_num=temporal_scale_num, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + self.decoder = Decoder3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + temporal_up_num=temporal_scale_num, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + self.use_slicing = True + + def encode(self, x: torch.FloatTensor, return_dict: bool = True): + h = self.slicing_encode(x) + posterior = DiagonalGaussianDistribution(h).mode() + + if not return_dict: + return (posterior,) + + return posterior + + def decode_( + self, z: torch.Tensor, return_dict: bool = True + ): + decoded = self.slicing_decode(z) + + if not return_dict: + return (decoded,) + + return decoded + + def _encode( + self, x, memory_state = MemoryState.DISABLED, memory_cache = None + ) -> torch.Tensor: + _x = x.to(self.device) + h = self.encoder(_x, memory_state=memory_state, memory_cache=memory_cache) + return h.to(x.device) + + def _decode( + self, z, memory_state = MemoryState.DISABLED, memory_cache = None + ) -> torch.Tensor: + _z = z.to(self.device) + output = self.decoder(_z, memory_state=memory_state, memory_cache=memory_cache) + return output.to(z.device) + + def slicing_encode(self, x: torch.Tensor) -> torch.Tensor: + if self.use_slicing and (x.shape[2] - 1) > self.slicing_sample_min_size: + memory_cache = {} + split_size = max( + self.slicing_sample_min_size, + getattr(self, "temporal_downsample_factor", 1), + ) + x_slices = list(x[:, :, 1:].split(split_size=split_size, dim=2)) + min_active_len = getattr(self, "temporal_downsample_factor", 1) + if len(x_slices) > 1 and x_slices[-1].shape[2] < min_active_len: + x_slices[-2] = torch.cat((x_slices[-2], x_slices[-1]), dim=2) + x_slices.pop() + encoded_slices = [ + self._encode( + torch.cat((x[:, :, :1], x_slices[0]), dim=2), + memory_state=MemoryState.INITIALIZING, + memory_cache=memory_cache, + ) + ] + for x_idx in range(1, len(x_slices)): + encoded_slices.append( + self._encode(x_slices[x_idx], memory_state=MemoryState.ACTIVE, memory_cache=memory_cache) + ) + out = torch.cat(encoded_slices, dim=2) + return out + else: + return self._encode(x) + + def slicing_decode(self, z: torch.Tensor) -> torch.Tensor: + if self.use_slicing and (z.shape[2] - 1) > self.slicing_latent_min_size: + memory_cache = {} + z_slices = z[:, :, 1:].split(split_size=self.slicing_latent_min_size, dim=2) + decoded_slices = [ + self._decode( + torch.cat((z[:, :, :1], z_slices[0]), dim=2), + memory_state=MemoryState.INITIALIZING, + memory_cache=memory_cache, + ) + ] + for z_idx in range(1, len(z_slices)): + decoded_slices.append( + self._decode(z_slices[z_idx], memory_state=MemoryState.ACTIVE, memory_cache=memory_cache) + ) + out = torch.cat(decoded_slices, dim=2) + return out + else: + return self._decode(z) + + def forward(self, x: torch.FloatTensor, mode: Literal["encode", "decode", "all"] = "all"): + def _unwrap(value): + return value[0] if isinstance(value, tuple) else value + + if mode == "encode": + return _unwrap(self.encode(x)) + if mode == "decode": + return _unwrap(self.decode_(x)) + if mode == "all": + latent = _unwrap(self.encode(x)) + return _unwrap(self.decode_(latent)) + raise ValueError(f"Unknown SeedVR2 VAE forward mode: {mode}") + +class VideoAutoencoderKLWrapper(VideoAutoencoderKL): + def __init__( + self, + spatial_downsample_factor = 8, + temporal_downsample_factor = 4, + ): + self.spatial_downsample_factor = spatial_downsample_factor + self.temporal_downsample_factor = temporal_downsample_factor + super().__init__() + self.set_memory_limit(BYTEDANCE_VAE_CONV_MEM_GIB, BYTEDANCE_VAE_NORM_MEM_GIB) + + def forward(self, x: torch.FloatTensor): + z, p = self._encode_with_raw_latent(x) + x = self.decode(z) + return x, z, p + + def _encode_with_raw_latent(self, x): + if x.ndim == 4: + x = x.unsqueeze(2) + x = x.to(dtype=next(self.parameters()).dtype) + self.device = x.device + p = super().encode(x) + z = p.squeeze(2) + return z, p + + def encode(self, x): + z, _ = self._encode_with_raw_latent(x) + return z + + def decode(self, z, seedvr2_tiling=None): + seedvr2_tiling = {} if seedvr2_tiling is None else seedvr2_tiling + if not isinstance(seedvr2_tiling, dict): + raise RuntimeError( + "SeedVR2 VideoAutoencoderKLWrapper.decode: `seedvr2_tiling` must be a dict; " + f"got {type(seedvr2_tiling).__name__} with value {seedvr2_tiling!r}." + ) + + if z.ndim == 5: + _, c, _, _, _ = z.shape + if c != SEEDVR2_LATENT_CHANNELS: + raise RuntimeError( + "SeedVR2 VideoAutoencoderKLWrapper.decode: 5-D latent input must " + f"have {SEEDVR2_LATENT_CHANNELS} channels; got shape {tuple(z.shape)}." + ) + latent = z + elif z.ndim == 4: + b, tc, h, w = z.shape + if tc % SEEDVR2_LATENT_CHANNELS != 0: + raise RuntimeError( + "SeedVR2 VideoAutoencoderKLWrapper.decode: 4-D latent input must " + f"use collapsed channel layout (B, {SEEDVR2_LATENT_CHANNELS}*T, H, W); " + f"got shape {tuple(z.shape)}." + ) + latent = z.reshape(b, SEEDVR2_LATENT_CHANNELS, -1, h, w) + else: + raise RuntimeError( + "SeedVR2 VideoAutoencoderKLWrapper.decode: latent input must be " + f"4-D collapsed (B, {SEEDVR2_LATENT_CHANNELS}*T, H, W) or " + f"5-D (B, {SEEDVR2_LATENT_CHANNELS}, T, H, W); " + f"got shape {tuple(z.shape)}." + ) + scale = BYTEDANCE_VAE_SCALING_FACTOR + shift = BYTEDANCE_VAE_SHIFTING_FACTOR + latent = latent / scale + shift + + self.device = latent.device + enable_tiling = seedvr2_tiling.get("enable_tiling", False) + + if enable_tiling: + decode_seedvr2_args = dict(seedvr2_tiling) + decode_seedvr2_args.pop("enable_tiling", None) + tile_h, tile_w = decode_seedvr2_args.get("tile_size", (512, 512)) + ov_h, ov_w = decode_seedvr2_args.get("tile_overlap", (64, 64)) + decode_seedvr2_args["tile_overlap"] = ( + min(ov_h, max(0, tile_h - 8)), + min(ov_w, max(0, tile_w - 8)), + ) + x = tiled_vae(latent, self, **decode_seedvr2_args, encode=False) + if x.ndim == 4: + # tiled_vae squeezes the temporal axis when + # temporal_downsample_factor == 1 AND latent T == 1 + # (see tiled_vae line 179-180); re-add it so the post-decode + # pipeline can keep batch and time distinct on the tiled path. + x = x.unsqueeze(2) + else: + x = super().decode_(latent) + + h, w = x.shape[-2:] + w2 = w - (w % 2) + h2 = h - (h % 2) + x = x[..., :h2, :w2] + + return x + + def decode_tiled(self, z, tile_x=32, tile_y=32, overlap=8, tile_t=None, overlap_t=None): + # SeedVR2's causal VAE owns temporal via the MemoryState cache; external + # temporal tiling breaks that continuity, so only spatial tiling is applied. + sf = self.spatial_downsample_factor + seedvr2_tiling = { + "enable_tiling": True, + "tile_size": (tile_y * sf, tile_x * sf), + "tile_overlap": (overlap * sf, overlap * sf), + "temporal_size": None, + "temporal_overlap": None, + } + return self.decode(z, seedvr2_tiling=seedvr2_tiling) + + def encode_tiled(self, x, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None): + # External temporal tiling knobs are discarded; the causal VAE keeps its + # own internal MemoryState slicing. + if tile_y is None: + tile_y = 512 + if tile_x is None: + tile_x = 512 + if overlap is None: + overlap_y = 64 + overlap_x = 64 + else: + overlap_y = overlap + overlap_x = overlap + overlap_y = min(overlap_y, max(0, tile_y - 8)) + overlap_x = min(overlap_x, max(0, tile_x - 8)) + self.device = x.device + return tiled_vae( + x, + self, + tile_size=(tile_y, tile_x), + tile_overlap=(overlap_y, overlap_x), + temporal_size=None, + temporal_overlap=None, + encode=True, + ) + + def comfy_format_encoded(self, samples): + if samples.ndim == 4: + samples = samples.unsqueeze(2) + samples = samples.contiguous() + samples = samples * BYTEDANCE_VAE_SCALING_FACTOR + return samples + + def comfy_memory_used_decode(self, shape): + bytes_per_output_pixel = 160 + + def output_pixels(latent_t, latent_h, latent_w): + output_t = max(1, (latent_t - 1) * 4 + 1) + return output_t * latent_h * 8 * latent_w * 8 + + # SeedVR2 decode performs full-frame LAB histogram matching: fp32 channels + # plus int64 sort indices dominate peak memory, not the VAE weight dtype. + if len(shape) == 5: + candidates = [] + if shape[1] == SEEDVR2_LATENT_CHANNELS: + candidates.append((shape[2], shape[3], shape[4])) + if shape[-1] == SEEDVR2_LATENT_CHANNELS: + candidates.append((shape[1], shape[2], shape[3])) + if len(candidates) == 0: + candidates.append((shape[2], shape[3], shape[4])) + pixels = max(output_pixels(*candidate) for candidate in candidates) + elif len(shape) == 4: + latent_t = max(1, (shape[1] + SEEDVR2_LATENT_CHANNELS - 1) // SEEDVR2_LATENT_CHANNELS) + pixels = output_pixels(latent_t, shape[2], shape[3]) + else: + pixels = output_pixels(1, shape[-2], shape[-1]) + return pixels * bytes_per_output_pixel + + def set_memory_limit(self, conv_max_mem: Optional[float], norm_max_mem: Optional[float]): + set_norm_limit(norm_max_mem) + for m in self.modules(): + if isinstance(m, InflatedCausalConv3d): + m.set_memory_limit(conv_max_mem if conv_max_mem is not None else float("inf")) diff --git a/comfy/model_base.py b/comfy/model_base.py index dcfa555dc..786a7c127 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -55,6 +55,7 @@ import comfy.ldm.pixeldit.model import comfy.ldm.pixeldit.pid import comfy.ldm.ace.model import comfy.ldm.omnigen.omnigen2 +import comfy.ldm.seedvr.model import comfy.ldm.boogu.model import comfy.ldm.qwen_image.model import comfy.ldm.ideogram4.model @@ -932,6 +933,17 @@ class HunyuanDiT(BaseModel): out['image_meta_size'] = comfy.conds.CONDRegular(torch.FloatTensor([[height, width, target_height, target_width, 0, 0]])) return out +class SeedVR2(BaseModel): + def __init__(self, model_config, model_type=ModelType.FLOW, device=None): + super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.seedvr.model.NaDiT) + + def extra_conds(self, **kwargs): + out = super().extra_conds(**kwargs) + condition = kwargs.get("condition", None) + if condition is not None: + out["condition"] = comfy.conds.CONDRegular(condition) + return out + class PixArt(BaseModel): def __init__(self, model_config, model_type=ModelType.EPS, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.pixart.pixartms.PixArtMS) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index e53d848c9..174bc77cc 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -598,6 +598,44 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): return dit_config + seedvr2_7b_separate_key = "{}blocks.35.mlp.vid.proj_out.weight".format(key_prefix) + if seedvr2_7b_separate_key in state_dict_keys and state_dict[seedvr2_7b_separate_key].shape[0] == 3072: # seedvr2 7b + dit_config = {} + dit_config["image_model"] = "seedvr2" + dit_config["vid_dim"] = 3072 + dit_config["heads"] = 24 + dit_config["num_layers"] = 36 + # This checkpoint uses separate vid/txt MMModule keys in every block. + dit_config["mm_layers"] = 36 + dit_config["norm_eps"] = 1e-5 + dit_config["rope_type"] = "rope3d" + dit_config["rope_dim"] = 64 + dit_config["mlp_type"] = "normal" + return dit_config + if "{}blocks.35.mlp.all.proj_in_gate.weight".format(key_prefix) in state_dict_keys: # seedvr2 7b + dit_config = {} + dit_config["image_model"] = "seedvr2" + dit_config["vid_dim"] = 3072 + dit_config["heads"] = 24 + dit_config["num_layers"] = 36 + # This checkpoint uses shared all.* MMModule keys after the initial blocks. + dit_config["mm_layers"] = 10 + dit_config["norm_eps"] = 1e-5 + dit_config["rope_type"] = "rope3d" + dit_config["rope_dim"] = 64 + dit_config["mlp_type"] = "swiglu" + return dit_config + if "{}blocks.31.mlp.all.proj_in_gate.weight".format(key_prefix) in state_dict_keys: # seedvr2 3b + dit_config = {} + dit_config["image_model"] = "seedvr2" + dit_config["vid_dim"] = 2560 + dit_config["heads"] = 20 + dit_config["num_layers"] = 32 + dit_config["norm_eps"] = 1.0e-05 + dit_config["mlp_type"] = "swiglu" + dit_config["vid_out_norm"] = True + return dit_config + if '{}head.modulation'.format(key_prefix) in state_dict_keys: # Wan 2.1 dit_config = {} dit_config["image_model"] = "wan2.1" @@ -1119,9 +1157,10 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): return unet_config -def model_config_from_unet_config(unet_config, state_dict=None): + +def model_config_from_unet_config(unet_config, state_dict=None, unet_key_prefix=""): for model_config in comfy.supported_models.models: - if model_config.matches(unet_config, state_dict): + if model_config.matches(unet_config, state_dict, unet_key_prefix=unet_key_prefix): return model_config(unet_config) logging.error("no match {}".format(unet_config)) @@ -1131,7 +1170,7 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal unet_config = detect_unet_config(state_dict, unet_key_prefix, metadata=metadata) if unet_config is None: return None - model_config = model_config_from_unet_config(unet_config, state_dict) + model_config = model_config_from_unet_config(unet_config, state_dict, unet_key_prefix) if model_config is None and use_base_if_no_match: model_config = comfy.supported_models_base.BASE(unet_config) diff --git a/comfy/sd.py b/comfy/sd.py index 071a3102a..4a0742e7a 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -16,6 +16,7 @@ import comfy.ldm.cosmos.vae import comfy.ldm.wan.vae import comfy.ldm.wan.vae2_2 import comfy.ldm.hunyuan3d.vae +import comfy.ldm.seedvr.vae import comfy.ldm.triposplat.vae import comfy.ldm.ace.vae.music_dcae_pipeline import comfy.ldm.cogvideo.vae @@ -473,7 +474,8 @@ class CLIP: class VAE: def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None): - if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format + is_seedvr2_vae = "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd + if not is_seedvr2_vae and 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format sd = diffusers_convert.convert_vae_state_dict(sd) if model_management.is_amd(): @@ -500,6 +502,8 @@ class VAE: self.upscale_index_formula = None self.extra_1d_channel = None self.crop_input = True + self.handles_tiling = False + self.format_encoded = None self.audio_sample_rate = 44100 @@ -546,6 +550,22 @@ class VAE: self.first_stage_model = StageC_coder() self.downscale_ratio = 32 self.latent_channels = 16 + elif "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd: # seedvr2 + self.first_stage_model = comfy.ldm.seedvr.vae.VideoAutoencoderKLWrapper() + self.latent_channels = comfy.ldm.seedvr.vae.SEEDVR2_LATENT_CHANNELS + self.latent_dim = 3 + self.disable_offload = True + self.memory_used_decode = lambda shape, dtype: self.first_stage_model.comfy_memory_used_decode(shape) + self.memory_used_encode = lambda shape, dtype: (max(shape[2], 5) * shape[3] * shape[4] * 64) * model_management.dtype_size(dtype) + self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32] + self.handles_tiling = True + self.format_encoded = self.first_stage_model.comfy_format_encoded + self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8) + self.downscale_index_formula = (4, 8, 8) + self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8) + self.upscale_index_formula = (4, 8, 8) + self.process_input = lambda image: image * 2.0 - 1.0 + self.crop_input = False elif "decoder.conv_in.weight" in sd: if sd['decoder.conv_in.weight'].shape[1] == 64: ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True} @@ -1012,6 +1032,10 @@ class VAE: decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype()) return self.process_output(comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, index_formulas=self.upscale_index_formula, output_device=self.output_device)) + def _decode_tiled_owned(self, samples, **kwargs): + out = self.first_stage_model.decode_tiled(samples.to(self.vae_dtype).to(self.device), **kwargs) + return self.process_output(out.to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True)) + def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64): steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap) steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap) @@ -1048,6 +1072,25 @@ class VAE: encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype()) return comfy.utils.tiled_scale_multidim(samples, encode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.downscale_ratio, out_channels=self.latent_channels, downscale=True, index_formulas=self.downscale_index_formula, output_device=self.output_device) + def _encode_tiled_owned(self, pixel_samples, **kwargs): + x = self.process_input(pixel_samples).to(self.vae_dtype).to(self.device) + out = self.first_stage_model.encode_tiled(x, **kwargs) + return out.to(device=self.output_device, dtype=self.vae_output_dtype()) + + def _owned_tiled_args(self, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None): + args = {} + if tile_x is not None: + args["tile_x"] = tile_x + if tile_y is not None: + args["tile_y"] = tile_y + if overlap is not None: + args["overlap"] = overlap + if tile_t is not None: + args["tile_t"] = tile_t + if overlap_t is not None: + args["overlap_t"] = overlap_t + return args + def decode(self, samples_in, vae_options={}): self.throw_exception_if_invalid() pixel_samples = None @@ -1095,11 +1138,19 @@ class VAE: if dims == 1 or self.extra_1d_channel is not None: pixel_samples = self.decode_tiled_1d(samples_in) elif dims == 2: - pixel_samples = self.decode_tiled_(samples_in) + if self.handles_tiling: + tile = 256 // self.spacial_compression_decode() + overlap = tile // 4 + pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap) + else: + pixel_samples = self.decode_tiled_(samples_in) elif dims == 3: tile = 256 // self.spacial_compression_decode() overlap = tile // 4 - pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) + if self.handles_tiling: + pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap) + else: + pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1) return pixel_samples @@ -1118,7 +1169,9 @@ class VAE: args["overlap"] = overlap with model_management.cuda_device_context(self.device): - if dims == 1 or self.extra_1d_channel is not None: + if self.handles_tiling and dims in (2, 3): + output = self._decode_tiled_owned(samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t)) + elif dims == 1 or self.extra_1d_channel is not None: args.pop("tile_y") output = self.decode_tiled_1d(samples, **args) elif dims == 2: @@ -1179,12 +1232,17 @@ class VAE: if self.latent_dim == 3: tile = 256 overlap = tile // 4 - samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) + if self.handles_tiling: + samples = self._encode_tiled_owned(pixel_samples, tile_x=tile, tile_y=tile, overlap=overlap) + else: + samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) elif self.latent_dim == 1 or self.extra_1d_channel is not None: samples = self.encode_tiled_1d(pixel_samples) else: samples = self.encode_tiled_(pixel_samples) + if self.format_encoded is not None: + samples = self.format_encoded(samples) return samples def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None): @@ -1192,7 +1250,7 @@ class VAE: pixel_samples = self.vae_encode_crop_pixels(pixel_samples) dims = self.latent_dim pixel_samples = pixel_samples.movedim(-1, 1) - if dims == 3: + if dims == 3 and pixel_samples.ndim < 5: if not self.not_video: pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0) else: @@ -1216,21 +1274,27 @@ class VAE: elif dims == 2: samples = self.encode_tiled_(pixel_samples, **args) elif dims == 3: - if tile_t is not None: - tile_t_latent = max(2, self.downscale_ratio[0](tile_t)) + if self.handles_tiling: + samples = self._encode_tiled_owned(pixel_samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t)) else: - tile_t_latent = 9999 - args["tile_t"] = self.upscale_ratio[0](tile_t_latent) + if tile_t is not None: + tile_t_latent = max(2, self.downscale_ratio[0](tile_t)) + else: + tile_t_latent = 9999 + args["tile_t"] = self.upscale_ratio[0](tile_t_latent) - if overlap_t is None: - args["overlap"] = (1, overlap, overlap) - else: - args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), overlap, overlap) - maximum = pixel_samples.shape[2] - maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum)) + spatial_overlap = overlap if overlap is not None else 64 + if overlap_t is None: + args["overlap"] = (1, spatial_overlap, spatial_overlap) + else: + args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), spatial_overlap, spatial_overlap) + maximum = pixel_samples.shape[2] + maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum)) - samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args) + samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args) + if self.format_encoded is not None: + samples = self.format_encoded(samples) return samples def get_sd(self): @@ -1898,7 +1962,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes) else: manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes) - model_config.set_inference_dtype(unet_dtype, manual_cast_dtype) + model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device) if model_config.clip_vision_prefix is not None: if output_clipvision: @@ -2039,7 +2103,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes) else: manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes) - model_config.set_inference_dtype(unet_dtype, manual_cast_dtype) + model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device) if custom_operations is not None: model_config.custom_operations = custom_operations diff --git a/comfy/supported_models.py b/comfy/supported_models.py index afb66e6f3..b82e4178f 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -1685,6 +1685,40 @@ class Chroma(supported_models_base.BASE): t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref)) return supported_models_base.ClipTarget(comfy.text_encoders.pixart_t5.PixArtTokenizer, comfy.text_encoders.pixart_t5.pixart_te(**t5_detect)) +class SeedVR2(supported_models_base.BASE): + unet_config = { + "image_model": "seedvr2" + } + unet_extra_config = {} + required_keys = { + "{}positive_conditioning", + "{}negative_conditioning", + } + latent_format = comfy.latent_formats.SeedVR2 + + vae_key_prefix = ["vae."] + text_encoder_key_prefix = ["text_encoders."] + supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32] + sampling_settings = { + "shift": 1.0, + } + + def set_inference_dtype(self, dtype, manual_cast_dtype, device=None): + if ( + dtype == torch.float16 + and manual_cast_dtype is None + and comfy.model_management.should_use_bf16(device) + ): + manual_cast_dtype = torch.bfloat16 + super().set_inference_dtype(dtype, manual_cast_dtype, device=device) + + def get_model(self, state_dict, prefix="", device=None): + out = model_base.SeedVR2(self, device=device) + return out + + def clip_target(self, state_dict={}): + return None + class ChromaRadiance(Chroma): unet_config = { "image_model": "chroma_radiance", @@ -2348,6 +2382,7 @@ models = [ HiDream, HiDreamO1, Chroma, + SeedVR2, ChromaRadiance, ACEStep, ACEStep15, diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py index 0e7a829ba..e3a8e131f 100644 --- a/comfy/supported_models_base.py +++ b/comfy/supported_models_base.py @@ -54,13 +54,13 @@ class BASE: optimizations = {"fp8": False} @classmethod - def matches(s, unet_config, state_dict=None): + def matches(s, unet_config, state_dict=None, unet_key_prefix=""): for k in s.unet_config: if k not in unet_config or s.unet_config[k] != unet_config[k]: return False if state_dict is not None: for k in s.required_keys: - if k not in state_dict: + if k.format(unet_key_prefix) not in state_dict: return False return True @@ -115,7 +115,7 @@ class BASE: replace_prefix = {"": self.vae_key_prefix[0]} return utils.state_dict_prefix_replace(state_dict, replace_prefix) - def set_inference_dtype(self, dtype, manual_cast_dtype): + def set_inference_dtype(self, dtype, manual_cast_dtype, device=None): self.unet_config['dtype'] = dtype self.manual_cast_dtype = manual_cast_dtype diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py index f050061ed..0bba8341b 100644 --- a/comfy/text_encoders/gemma4.py +++ b/comfy/text_encoders/gemma4.py @@ -1088,7 +1088,7 @@ class Gemma4_Tokenizer(): h, w = samples.shape[2], samples.shape[3] patch_size = 16 pooling_k = 3 - max_soft_tokens = 70 if is_video else 280 # video uses smaller token budget per frame + max_soft_tokens = kwargs.get("max_soft_tokens", 70 if is_video else 280) max_patches = max_soft_tokens * pooling_k * pooling_k target_px = max_patches * patch_size * patch_size factor = (target_px / (h * w)) ** 0.5 diff --git a/comfy_extras/nodes_seedvr.py b/comfy_extras/nodes_seedvr.py new file mode 100644 index 000000000..c4ca3b55c --- /dev/null +++ b/comfy_extras/nodes_seedvr.py @@ -0,0 +1,614 @@ +import logging + +from typing_extensions import override +from comfy_api.latest import ComfyExtension, io +import torch + +import comfy.model_management +from comfy.ldm.seedvr.color_fix import ( + adain_color_transfer, + lab_color_transfer, + wavelet_color_transfer, +) +from comfy.ldm.seedvr.constants import ( + BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE, + SEEDVR2_ADAIN_SCALE_MULTIPLIER, + SEEDVR2_CHUNK_GIB_PER_MPX_FRAME, + SEEDVR2_CHUNK_RESERVED_GIB, + SEEDVR2_CHUNK_SIGMA_GIB, + SEEDVR2_CHUNK_SIGMA_K, + SEEDVR2_COLOR_MEM_HEADROOM, + SEEDVR2_DTYPE_BYTES_FLOOR, + SEEDVR2_LAB_SCALE_MULTIPLIER, + SEEDVR2_LATENT_CHANNELS, + SEEDVR2_OOM_BACKOFF_DIVISOR, + SEEDVR2_WAVELET_SCALE_MULTIPLIER, +) + +from torchvision.transforms import functional as TVF +from torchvision.transforms.functional import InterpolationMode + + +_SEEDVR2_INVALID_MODEL_MSG_PREFIX = "SeedVR2Conditioning: model object does not match expected SeedVR2 structure" +_ATTR_MISSING = object() + + +def _resolve_seedvr2_diffusion_model(model): + inner = getattr(model, "model", _ATTR_MISSING) + if inner is _ATTR_MISSING: + raise RuntimeError( + f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: input has no 'model' attribute " + f"(got type {type(model).__name__})." + ) + if inner is None: + raise RuntimeError( + f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: input.model is None " + f"(input type {type(model).__name__})." + ) + diffusion_model = getattr(inner, "diffusion_model", _ATTR_MISSING) + if diffusion_model is _ATTR_MISSING: + raise RuntimeError( + f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: 'model.model' has no " + f"'diffusion_model' attribute (got type {type(inner).__name__})." + ) + if diffusion_model is None: + raise RuntimeError( + f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: 'model.model.diffusion_model' " + f"is None (model.model type {type(inner).__name__})." + ) + return diffusion_model + + +def div_pad(image, factor): + height_factor, width_factor = factor + height, width = image.shape[-2:] + + pad_height = (height_factor - (height % height_factor)) % height_factor + pad_width = (width_factor - (width % width_factor)) % width_factor + + if pad_height == 0 and pad_width == 0: + return image + + padding = (0, pad_width, 0, pad_height) + return torch.nn.functional.pad(image, padding, mode='constant', value=0.0) + +def cut_videos(videos): + t = videos.size(1) + if t < 1: + raise ValueError("SeedVR2Preprocess expected at least one frame.") + if t == 1: + return videos + if t <= 4: + padding = videos[:, -1:].repeat(1, 4 - t + 1, 1, 1, 1) + return torch.cat([videos, padding], dim=1) + if (t - 1) % 4 == 0: + return videos + padding = videos[:, -1:].repeat(1, 4 - ((t - 1) % 4), 1, 1, 1) + videos = torch.cat([videos, padding], dim=1) + if (videos.size(1) - 1) % 4 != 0: + raise ValueError(f"SeedVR2Preprocess failed to pad video length to 4n+1; got {videos.size(1)} frames.") + return videos + +def _seedvr2_input_shorter_edge(images, node_name): + if images.dim() == 4: + return min(images.shape[1], images.shape[2]) + if images.dim() == 5: + return min(images.shape[2], images.shape[3]) + raise ValueError( + f"{node_name}: expected 4-D or 5-D IMAGE tensor, " + f"got shape {tuple(images.shape)}" + ) + + +def _seedvr2_pad(images, upscaled_shorter_edge, node_name): + if upscaled_shorter_edge < 2: + raise ValueError( + f"{node_name}: input shorter edge must be at least 2 pixels; " + f"got {upscaled_shorter_edge}." + ) + if images.shape[-1] > 3: + images = images[..., :3] + if images.dim() == 4: + # Comfy video components arrive as a 4-D IMAGE frame sequence: + # (frames, H, W, C). SeedVR2 consumes that as one video. + images = images.unsqueeze(0) + elif images.dim() != 5: + raise ValueError( + f"{node_name}: expected 4-D or 5-D IMAGE tensor, " + f"got shape {tuple(images.shape)}" + ) + images = images.permute(0, 1, 4, 2, 3) + + b, t, c, h, w = images.shape + images = images.reshape(b * t, c, h, w) + + images = torch.clamp(images, 0.0, 1.0) + images = div_pad(images, (16, 16)) + _, _, new_h, new_w = images.shape + + images = images.reshape(b, t, c, new_h, new_w) + images = cut_videos(images) + images_bthwc = images.permute(0, 1, 3, 4, 2).contiguous() + + return io.NodeOutput(images_bthwc) + + +class SeedVR2Preprocess(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SeedVR2Preprocess", + display_name="Pre-Process SeedVR2 Input", + category="image/pre-processors", + description="Pad a resized image for SeedVR2 model. Alpha channel is dropped. The node Post-Process SeedVR2 Output re-applies it from the original resized image.", + search_aliases=["seedvr2", "upscale", "video upscale", "pad", "preprocess"], + inputs=[ + io.Image.Input("resized_images", tooltip="The resized image to process."), + ], + outputs=[ + io.Image.Output("images", tooltip="The padded image for VAE encoding."), + ] + ) + + @classmethod + def execute(cls, resized_images): + upscaled_shorter_edge = _seedvr2_input_shorter_edge(resized_images, "SeedVR2Preprocess") + return _seedvr2_pad( + resized_images, upscaled_shorter_edge, "SeedVR2Preprocess", + ) + + +class SeedVR2PostProcessing(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SeedVR2PostProcessing", + display_name="Post-Process SeedVR2 Output", + category="image/post-processors", + description="Align the generated image with the original resized image and apply color correction.", + search_aliases=["seedvr2", "upscale", "color correction", "color match", "postprocess"], + inputs=[ + io.Image.Input("images", tooltip="The generated image to process."), + io.Image.Input("original_resized_images", tooltip="The original resized image before pre-processing, used as reference."), + io.Combo.Input("color_correction_method", options=["lab", "wavelet", "adain", "none"], default="lab", tooltip="Method to match the generated image colors to the original image. lab: transfer color in CIELAB space, preserving detail (most faithful). wavelet: transfer low-frequency color, keeping upscaled high-frequency detail. adain: match per-channel mean/std (fastest, global tint). none: skip color transfer (geometry alignment only)."), + ], + outputs=[io.Image.Output(display_name="images", tooltip="The aligned, color-corrected image.")], + ) + + @classmethod + def execute(cls, images, original_resized_images, color_correction_method): + alpha_input = None + if original_resized_images.shape[-1] == 4: + alpha_input = original_resized_images[..., 3:4] + original_resized_images = original_resized_images[..., :3] + decoded_5d, decoded_was_4d = cls._as_bthwc(images) + reference_full, _ = cls._as_bthwc(original_resized_images) + decoded_5d = cls._restore_reference_batch_time(decoded_5d, reference_full) + + b = min(decoded_5d.shape[0], reference_full.shape[0]) + t = min(decoded_5d.shape[1], reference_full.shape[1]) + reference_h = reference_full.shape[2] + reference_w = reference_full.shape[3] + + decoded_5d = decoded_5d[:b, :t, :, :, :] + target_h = min(decoded_5d.shape[2], reference_h) + target_w = min(decoded_5d.shape[3], reference_w) + decoded_5d = decoded_5d[:, :, :target_h, :target_w, :] + if color_correction_method in ("lab", "wavelet", "adain"): + reference_5d = reference_full[:b, :t, :, :, :] + reference_5d = cls._resize_reference(reference_5d, target_h, target_w) + output_device = decoded_5d.device + decoded_raw = cls._to_seedvr2_raw(decoded_5d) + reference_raw = cls._to_seedvr2_raw(reference_5d) + decoded_flat = decoded_raw.permute(0, 1, 4, 2, 3).reshape(b * t, decoded_raw.shape[4], target_h, target_w) + reference_flat = reference_raw.permute(0, 1, 4, 2, 3).reshape(b * t, reference_raw.shape[4], target_h, target_w) + output = cls._color_transfer_chunked( + decoded_flat, reference_flat, output_device, color_correction_method, + ) + output = output.reshape(b, t, output.shape[1], output.shape[2], output.shape[3]).permute(0, 1, 3, 4, 2) + output = output.add(1.0).div(2.0).clamp(0.0, 1.0) + elif color_correction_method == "none": + output = decoded_5d + else: + raise ValueError(f"SeedVR2PostProcessing: unknown color_correction_method {color_correction_method!r}") + + if alpha_input is not None: + alpha_5d, _ = cls._as_bthwc(alpha_input) + alpha_5d = alpha_5d[:output.shape[0], :output.shape[1], :output.shape[2], :output.shape[3], :] + output = torch.cat([output, alpha_5d.to(dtype=output.dtype, device=output.device)], dim=-1) + h2 = output.shape[-3] - (output.shape[-3] % 2) + w2 = output.shape[-2] - (output.shape[-2] % 2) + output = output[:, :, :h2, :w2, :] + if decoded_was_4d: + output = output.reshape(-1, output.shape[-3], output.shape[-2], output.shape[-1]) + return io.NodeOutput(output) + + @staticmethod + def _as_bthwc(images): + if images.ndim == 4: + return images.unsqueeze(0), True + if images.ndim == 5: + return images, False + raise ValueError( + f"SeedVR2PostProcessing: expected 4-D or 5-D IMAGE tensor, got shape {tuple(images.shape)}" + ) + + @staticmethod + def _restore_reference_batch_time(decoded, reference): + if decoded.shape[0] != 1: + return decoded + ref_b, ref_t = reference.shape[:2] + if ref_b < 1 or decoded.shape[1] % ref_b != 0: + return decoded + decoded_t = decoded.shape[1] // ref_b + if decoded_t < ref_t: + return decoded + return decoded.reshape(ref_b, decoded_t, decoded.shape[2], decoded.shape[3], decoded.shape[4]) + + @staticmethod + def _to_seedvr2_raw(images): + return images.mul(2.0).sub(1.0) + + @staticmethod + def _color_transfer_on_vae_device(decoded_flat, reference_flat, output_device, transfer_fn): + color_device = comfy.model_management.vae_device() + decoded_flat = decoded_flat.to(device=color_device) + reference_flat = reference_flat.to(device=color_device) + output = transfer_fn(decoded_flat, reference_flat) + return output.to(device=output_device) + + @staticmethod + def _lab_color_transfer_on_vae_device(decoded_flat, reference_flat, output_device): + color_device = comfy.model_management.vae_device() + result = None + for start in range(decoded_flat.shape[0]): + decoded_frame = decoded_flat[start:start + 1].to(device=color_device).clone() + reference_frame = reference_flat[start:start + 1].to(device=color_device).clone() + output = lab_color_transfer(decoded_frame, reference_frame).to(device=output_device) + if result is None: + result = torch.empty( + (decoded_flat.shape[0],) + tuple(output.shape[1:]), + device=output_device, + dtype=output.dtype, + ) + result[start:start + 1].copy_(output) + if result is None: + raise ValueError("SeedVR2PostProcessing: LAB color correction requires at least one frame.") + return result + + @classmethod + def _color_transfer_chunked(cls, decoded_flat, reference_flat, output_device, color_correction_method): + chunk_size = cls._estimate_color_correction_chunk_size(decoded_flat, color_correction_method) + while True: + try: + return cls._run_color_transfer_chunks( + decoded_flat, reference_flat, output_device, color_correction_method, chunk_size, + ) + except Exception as e: + comfy.model_management.raise_non_oom(e) + if chunk_size <= 1: + raise RuntimeError( + "SeedVR2PostProcessing: color correction OOM at one frame; " + f"color_correction_method={color_correction_method}, shape={tuple(decoded_flat.shape)}." + ) from e + chunk_size = max(1, chunk_size // SEEDVR2_OOM_BACKOFF_DIVISOR) + + @classmethod + def _run_color_transfer_chunks(cls, decoded_flat, reference_flat, output_device, color_correction_method, chunk_size): + result = None + for start in range(0, decoded_flat.shape[0], chunk_size): + end = min(start + chunk_size, decoded_flat.shape[0]) + decoded_chunk = decoded_flat[start:end] + reference_chunk = reference_flat[start:end] + if color_correction_method == "lab": + output = cls._lab_color_transfer_on_vae_device(decoded_chunk, reference_chunk, output_device) + elif color_correction_method == "wavelet": + output = cls._color_transfer_on_vae_device( + decoded_chunk, reference_chunk, output_device, wavelet_color_transfer, + ) + else: + output = cls._color_transfer_on_vae_device( + decoded_chunk, reference_chunk, output_device, adain_color_transfer, + ) + if result is None: + result = torch.empty( + (decoded_flat.shape[0],) + tuple(output.shape[1:]), + device=output_device, + dtype=output.dtype, + ) + result[start:end].copy_(output) + if result is None: + raise ValueError("SeedVR2PostProcessing: color correction requires at least one frame.") + return result + + @classmethod + def _estimate_color_correction_chunk_size(cls, decoded_flat, color_correction_method): + multiplier = cls._color_correction_memory_multiplier(color_correction_method) + frames = decoded_flat.shape[0] + _, channels, height, width = decoded_flat.shape + dtype_bytes = max(decoded_flat.element_size(), SEEDVR2_DTYPE_BYTES_FLOOR) + bytes_per_frame = height * width * channels * dtype_bytes * multiplier + if bytes_per_frame <= 0: + return frames + color_device = comfy.model_management.vae_device() + free_memory = comfy.model_management.get_free_memory(color_device) + chunk_size = int((free_memory * SEEDVR2_COLOR_MEM_HEADROOM) // bytes_per_frame) + return max(1, min(frames, chunk_size)) + + @staticmethod + def _color_correction_memory_multiplier(color_correction_method): + if color_correction_method == "lab": + return SEEDVR2_LAB_SCALE_MULTIPLIER + if color_correction_method == "wavelet": + return SEEDVR2_WAVELET_SCALE_MULTIPLIER + if color_correction_method == "adain": + return SEEDVR2_ADAIN_SCALE_MULTIPLIER + raise ValueError(f"SeedVR2PostProcessing: unknown color_correction_method {color_correction_method!r}") + + @staticmethod + def _resize_reference(reference, height, width): + if reference.shape[2] == height and reference.shape[3] == width: + return reference + b, t = reference.shape[:2] + reference_flat = reference.permute(0, 1, 4, 2, 3).reshape(b * t, reference.shape[4], reference.shape[2], reference.shape[3]) + resized = TVF.resize( + reference_flat, + size=(height, width), + interpolation=InterpolationMode.BICUBIC, + antialias=not (isinstance(reference_flat, torch.Tensor) and reference_flat.device.type == "mps"), + ) + return resized.reshape(b, t, resized.shape[1], height, width).permute(0, 1, 3, 4, 2) + + +class SeedVR2Conditioning(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SeedVR2Conditioning", + display_name="Apply SeedVR2 Conditioning", + category="model/conditioning", + description="Build SeedVR2 positive/negative conditioning from a VAE latent.", + search_aliases=["seedvr2", "upscale", "conditioning"], + inputs=[ + io.Model.Input("model", tooltip="The SeedVR2 model."), + io.Latent.Input("vae_conditioning", display_name="latent"), + ], + outputs=[ + io.Conditioning.Output(display_name="positive", tooltip="The positive conditioning for sampling."), + io.Conditioning.Output(display_name="negative", tooltip="The negative conditioning for sampling."), + ], + ) + + @classmethod + def execute(cls, model, vae_conditioning) -> io.NodeOutput: + + vae_conditioning = vae_conditioning["samples"] + if vae_conditioning.ndim != 5: + raise ValueError( + "SeedVR2Conditioning expects a 5-D VAE latent in Comfy " + f"channel-first layout; got shape {tuple(vae_conditioning.shape)}." + ) + if vae_conditioning.shape[1] != SEEDVR2_LATENT_CHANNELS: + if vae_conditioning.shape[-1] == SEEDVR2_LATENT_CHANNELS: + raise ValueError( + "SeedVR2Conditioning expects SeedVR2 VAE latents in Comfy " + f"channel-first layout (B, {SEEDVR2_LATENT_CHANNELS}, T, H, W); " + f"got channel-last shape {tuple(vae_conditioning.shape)}." + ) + raise ValueError( + "SeedVR2Conditioning expects SeedVR2 VAE latents with " + f"{SEEDVR2_LATENT_CHANNELS} channels; got shape {tuple(vae_conditioning.shape)}." + ) + vae_conditioning = vae_conditioning.movedim(1, -1).contiguous() + model = _resolve_seedvr2_diffusion_model(model) + pos_cond = model.positive_conditioning + neg_cond = model.negative_conditioning + + mask = vae_conditioning.new_ones(vae_conditioning.shape[:-1] + (1,)) + condition = torch.cat((vae_conditioning, mask), dim=-1) + condition = condition.movedim(-1, 1) + + negative = [[neg_cond.unsqueeze(0), {"condition": condition}]] + positive = [[pos_cond.unsqueeze(0), {"condition": condition}]] + + return io.NodeOutput(positive, negative) + +def _seedvr2_chunk_crossfade_weights(overlap, device, dtype): + """Descending previous-chunk weights across the overlap (next chunk gets ``1 - w``): a Hann fade over the middle third, flat shoulders on the outer thirds.""" + ramp = torch.linspace(0.0, 1.0, steps=overlap, device=device, dtype=dtype) + ramp = ((ramp - 1.0 / 3.0) / (1.0 / 3.0)).clamp(0.0, 1.0) + return 0.5 + 0.5 * torch.cos(torch.pi * ramp) + + +class SeedVR2TemporalChunk(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SeedVR2TemporalChunk", + display_name="Split SeedVR2 Latent", + category="model/latent/batch", + description="Split a SeedVR2 video latent into overlapping temporal chunks small enough to sample one at a time within VRAM, wiring latents outputs to both Apply SeedVR2 Conditioning and the sampler latent input before recombining with Merge SeedVR2 Latents.", + search_aliases=["seedvr2", "split", "chunk", "temporal", "video upscale", "rebatch"], + inputs=[ + io.Latent.Input("latent", tooltip="The VAE-encoded SeedVR2 latent to split."), + io.Int.Input("temporal_overlap", default=0, min=0, max=16384, + tooltip="Latent frames shared between adjacent chunks and crossfaded at merge; 0 = no overlap."), + io.DynamicCombo.Input("chunking_mode", + tooltip="manual = use frames_per_chunk exactly; auto = predict the largest chunk that fits free VRAM.", + options=[ + io.DynamicCombo.Option("auto", []), + io.DynamicCombo.Option("manual", [ + io.Int.Input("frames_per_chunk", default=21, min=1, max=16384, step=4, + tooltip="Pixel frames per temporal chunk (4n+1: 1, 5, 9, 13, ...)."), + ]), + ]), + ], + outputs=[ + io.Latent.Output(display_name="latents", is_output_list=True, + tooltip="The temporal chunks in sequence order."), + io.Int.Output(display_name="temporal_overlap", + tooltip="The effective latent-frame overlap between adjacent chunks, for Merge SeedVR2 Latents."), + ], + ) + + @classmethod + def execute(cls, latent, temporal_overlap, chunking_mode) -> io.NodeOutput: + samples = latent["samples"] + if samples.ndim != 5: + raise ValueError( + f"SeedVR2TemporalChunk: expected a 5-D video latent (B, C, T, H, W); " + f"got shape {tuple(samples.shape)}." + ) + if samples.shape[1] != SEEDVR2_LATENT_CHANNELS: + raise ValueError( + f"SeedVR2TemporalChunk: expected {SEEDVR2_LATENT_CHANNELS} latent channels; " + f"got shape {tuple(samples.shape)}." + ) + if temporal_overlap < 0: + raise ValueError( + f"SeedVR2TemporalChunk: temporal_overlap must be >= 0; got {temporal_overlap}." + ) + mode = chunking_mode["chunking_mode"] + if mode not in ("auto", "manual"): + raise ValueError( + f"SeedVR2TemporalChunk: chunking_mode must be 'auto' or 'manual'; " + f"got {mode!r}." + ) + t_latent = samples.shape[2] + t_pixel = 4 * (t_latent - 1) + 1 + + if mode == "auto": + free_gb = comfy.model_management.get_free_memory( + comfy.model_management.get_torch_device()) / (1024 ** 3) + mpx_per_frame = (samples.shape[0] * samples.shape[3] * samples.shape[4]) * (BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE ** 2) / 1e6 + budget_gb = free_gb - SEEDVR2_CHUNK_RESERVED_GIB - SEEDVR2_CHUNK_SIGMA_K * SEEDVR2_CHUNK_SIGMA_GIB + chunk_latent_max = max(1, int(budget_gb / (SEEDVR2_CHUNK_GIB_PER_MPX_FRAME * mpx_per_frame))) + frames_per_chunk = min(4 * (chunk_latent_max - 1) + 1, t_pixel) + logging.info( + "SeedVR2TemporalChunk auto: free=%.2fGiB, %.2fMpx -> frames_per_chunk=%d (t_pixel=%d).", + free_gb, mpx_per_frame, frames_per_chunk, t_pixel, + ) + else: + frames_per_chunk = chunking_mode["frames_per_chunk"] + if frames_per_chunk < 1 or (frames_per_chunk - 1) % 4 != 0: + raise ValueError( + f"SeedVR2TemporalChunk: frames_per_chunk must be a 4n+1 pixel-frame count " + f"(1, 5, 9, 13, 17, 21, ...); got {frames_per_chunk}." + ) + + if t_pixel <= frames_per_chunk: + return io.NodeOutput([latent], 0) + + chunk_latent = (frames_per_chunk - 1) // 4 + 1 + temporal_overlap = min(temporal_overlap, chunk_latent - 1) + step = chunk_latent - temporal_overlap + + chunks = [] + for start in range(0, t_latent, step): + end = min(start + chunk_latent, t_latent) + chunk = latent.copy() + chunk["samples"] = samples[:, :, start:end].contiguous() + chunks.append(chunk) + if end >= t_latent: + break + return io.NodeOutput(chunks, temporal_overlap) + + +class SeedVR2TemporalMerge(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SeedVR2TemporalMerge", + display_name="Merge SeedVR2 Latents", + category="model/latent/batch", + is_input_list=True, + description="Recombine sampled SeedVR2 latent temporal chunks into one latent, crossfading each overlap with a Hann window sized by the temporal_overlap wired from Split SeedVR2 Latent.", + search_aliases=["seedvr2", "merge", "temporal", "hann", "crossfade"], + inputs=[ + io.Latent.Input("latents", tooltip="The sampled temporal chunks in sequence order."), + io.Int.Input("temporal_overlap", default=0, min=0, max=16384, force_input=True, + tooltip="The temporal_overlap output of Split SeedVR2 Latent. 0 = plain concatenation."), + ], + outputs=[ + io.Latent.Output(display_name="latent", tooltip="The recombined full-length latent."), + ], + ) + + @classmethod + def execute(cls, latents, temporal_overlap) -> io.NodeOutput: + temporal_overlap = temporal_overlap[0] + if temporal_overlap < 0: + raise ValueError( + f"SeedVR2TemporalMerge: temporal_overlap must be >= 0; got {temporal_overlap}." + ) + chunks = [entry["samples"] for entry in latents] + first = chunks[0] + if first.ndim != 5: + raise ValueError( + f"SeedVR2TemporalMerge: expected 5-D video latents (B, C, T, H, W); " + f"chunk 0 has shape {tuple(first.shape)}." + ) + for i, chunk in enumerate(chunks[1:], start=1): + if chunk.shape[:2] != first.shape[:2] or chunk.shape[3:] != first.shape[3:]: + raise ValueError( + f"SeedVR2TemporalMerge: chunk {i} shape {tuple(chunk.shape)} does not " + f"match chunk 0 shape {tuple(first.shape)} outside the temporal axis." + ) + if i < len(chunks) - 1 and chunk.shape[2] != first.shape[2]: + raise ValueError( + f"SeedVR2TemporalMerge: chunk {i} has {chunk.shape[2]} latent frames but " + f"chunk 0 has {first.shape[2]}; only the final chunk may be shorter." + ) + + out = latents[0].copy() + out.pop("noise_mask", None) + + if len(chunks) == 1: + out["samples"] = first + return io.NodeOutput(out) + if temporal_overlap == 0: + out["samples"] = torch.cat(chunks, dim=2) + return io.NodeOutput(out) + + chunk_latent = first.shape[2] + step = chunk_latent - min(temporal_overlap, chunk_latent - 1) + t_total = step * (len(chunks) - 1) + chunks[-1].shape[2] + b, c, _, h, w = first.shape + merged = torch.empty((b, c, t_total, h, w), device=first.device, dtype=first.dtype) + + merged[:, :, :chunk_latent] = first + filled = chunk_latent + for i, chunk in enumerate(chunks[1:], start=1): + start = i * step + end = start + chunk.shape[2] + # Crossfade width is bounded by the previous fill frontier and by a runt + # final chunk shorter than the configured overlap. + fade = min(filled - start, chunk.shape[2]) + if fade > 0: + w_prev = _seedvr2_chunk_crossfade_weights( + fade, chunk.device, chunk.dtype).view(1, 1, fade, 1, 1) + merged[:, :, start:start + fade] = ( + merged[:, :, start:start + fade] * w_prev + chunk[:, :, :fade] * (1.0 - w_prev) + ) + merged[:, :, start + fade:end] = chunk[:, :, fade:] + else: + merged[:, :, start:end] = chunk + filled = end + + out["samples"] = merged + return io.NodeOutput(out) + + +class SeedVRExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + SeedVR2Conditioning, + SeedVR2Preprocess, + SeedVR2PostProcessing, + SeedVR2TemporalChunk, + SeedVR2TemporalMerge, + ] + +async def comfy_entrypoint() -> SeedVRExtension: + return SeedVRExtension() diff --git a/nodes.py b/nodes.py index e126576fe..474e188fe 100644 --- a/nodes.py +++ b/nodes.py @@ -2458,6 +2458,7 @@ async def init_builtin_extra_nodes(): "nodes_camera_trajectory.py", "nodes_edit_model.py", "nodes_tcfg.py", + "nodes_seedvr.py", "nodes_context_windows.py", "nodes_qwen.py", "nodes_boogu.py", diff --git a/tests-unit/comfy_extras_test/test_seedvr2_conditioning.py b/tests-unit/comfy_extras_test/test_seedvr2_conditioning.py new file mode 100644 index 000000000..045502b5b --- /dev/null +++ b/tests-unit/comfy_extras_test/test_seedvr2_conditioning.py @@ -0,0 +1,186 @@ +"""SeedVR2 conditioning node regression tests.""" + +import importlib +import sys +from unittest.mock import MagicMock + +import pytest +import torch +import torch.nn as nn + +from comfy.cli_args import args as cli_args +from comfy.ldm.seedvr.constants import SEEDVR2_LATENT_CHANNELS + +if not torch.cuda.is_available(): + cli_args.cpu = True + + +_SENTINEL = object() +_TARGETS = ( + ("comfy.model_management", "comfy"), + ("comfy_extras.nodes_seedvr", "comfy_extras"), +) + + +def _import_nodes_seedvr_isolated(): + """Import comfy_extras.nodes_seedvr with comfy.model_management mocked.""" + priors = [] + for mod_name, parent_name in _TARGETS: + prior_mod = sys.modules.get(mod_name, _SENTINEL) + parent = sys.modules.get(parent_name) + attr = mod_name.split(".")[-1] + prior_attr = ( + getattr(parent, attr, _SENTINEL) if parent is not None else _SENTINEL + ) + priors.append((mod_name, parent_name, attr, prior_mod, prior_attr)) + + mock_mm = MagicMock() + for fn in ( + "xformers_enabled", "xformers_enabled_vae", + "pytorch_attention_enabled", "pytorch_attention_enabled_vae", + "sage_attention_enabled", "flash_attention_enabled", + "is_intel_xpu", + ): + getattr(mock_mm, fn).return_value = False + tv = torch.version.__version__.split(".") + mock_mm.torch_version_numeric = (int(tv[0]), int(tv[1])) + mock_mm.WINDOWS = False + sys.modules["comfy.model_management"] = mock_mm + if sys.modules.get("comfy") is None: + importlib.import_module("comfy") + comfy_pkg = sys.modules.get("comfy") + if comfy_pkg is not None: + setattr(comfy_pkg, "model_management", mock_mm) + nodes_seedvr = sys.modules.get("comfy_extras.nodes_seedvr") or ( + importlib.import_module("comfy_extras.nodes_seedvr") + ) + + def _restore(): + for mod_name, parent_name, attr, prior_mod, prior_attr in priors: + if prior_mod is _SENTINEL: + sys.modules.pop(mod_name, None) + else: + sys.modules[mod_name] = prior_mod + parent = sys.modules.get(parent_name) + if parent is None: + continue + if prior_attr is _SENTINEL: + if hasattr(parent, attr): + delattr(parent, attr) + else: + setattr(parent, attr, prior_attr) + + return nodes_seedvr, _restore + + +class _Rope(nn.Module): + def __init__(self): + super().__init__() + self.freqs = nn.Parameter(torch.zeros(4)) + + +class _Block(nn.Module): + def __init__(self): + super().__init__() + self.rope = _Rope() + + +class _DiffusionModel(nn.Module): + def __init__(self, n_blocks=3, conditioning_dtype=torch.float32): + super().__init__() + self.blocks = nn.ModuleList([_Block() for _ in range(n_blocks)]) + self.register_buffer("positive_conditioning", torch.ones((2, 4), dtype=conditioning_dtype)) + self.register_buffer("negative_conditioning", torch.zeros((3, 4), dtype=conditioning_dtype)) + + +class _ModelInner: + def __init__(self, diffusion_model): + self.diffusion_model = diffusion_model + + +class _ModelPatcher: + def __init__(self, diffusion_model): + self.model = _ModelInner(diffusion_model) + + +def test_seedvr2_conditioning_schema_exposes_conditioning_outputs(): + nodes_seedvr, restore = _import_nodes_seedvr_isolated() + try: + schema = nodes_seedvr.SeedVR2Conditioning.define_schema() + assert [input_item.id for input_item in schema.inputs] == [ + "model", + "vae_conditioning", + ] + assert schema.inputs[1].display_name == "latent" + assert [output.display_name for output in schema.outputs] == [ + "positive", + "negative", + ] + finally: + restore() + + +def test_seedvr2_conditioning_rejects_wrong_latent_channels(): + nodes_seedvr, restore = _import_nodes_seedvr_isolated() + try: + patcher = _ModelPatcher(_DiffusionModel()) + vae_conditioning = {"samples": torch.zeros(1, 8, 2, 2, 2)} + + with pytest.raises(ValueError, match=f"{SEEDVR2_LATENT_CHANNELS} channels"): + nodes_seedvr.SeedVR2Conditioning.execute(patcher, vae_conditioning) + finally: + restore() + + +def test_seedvr2_conditioning_returns_conditioning_deterministically(): + nodes_seedvr, restore = _import_nodes_seedvr_isolated() + try: + diffusion_model = _DiffusionModel() + patcher = _ModelPatcher(diffusion_model) + samples = torch.arange( + 1, + 1 + SEEDVR2_LATENT_CHANNELS * 3 * 2 * 2, + dtype=torch.float32, + ).reshape(1, SEEDVR2_LATENT_CHANNELS, 3, 2, 2) + vae_conditioning = {"samples": samples} + + first_positive, first_negative = ( + nodes_seedvr.SeedVR2Conditioning.execute( + patcher, + vae_conditioning, + ) + ) + second_positive, second_negative = ( + nodes_seedvr.SeedVR2Conditioning.execute( + patcher, + vae_conditioning, + ) + ) + + channel_last = samples.movedim(1, -1).contiguous() + expected_condition = torch.cat( + [ + channel_last, + torch.ones((*channel_last.shape[:-1], 1)), + ], + dim=-1, + ).movedim(-1, 1) + + assert torch.equal( + first_positive[0][1]["condition"], + expected_condition, + ) + assert torch.equal( + second_positive[0][1]["condition"], + expected_condition, + ) + assert torch.equal( + first_negative[0][1]["condition"], + expected_condition, + ) + assert torch.equal( + second_negative[0][1]["condition"], + expected_condition, + ) + finally: + restore() diff --git a/tests-unit/comfy_extras_test/test_seedvr2_nodes.py b/tests-unit/comfy_extras_test/test_seedvr2_nodes.py new file mode 100644 index 000000000..1c5d20ac9 --- /dev/null +++ b/tests-unit/comfy_extras_test/test_seedvr2_nodes.py @@ -0,0 +1,55 @@ +import importlib +import inspect +import sys +from unittest.mock import MagicMock, patch + +import torch + +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + + +def test_seedvr_node_signature_matches_schema(): + mock_mm = MagicMock() + mock_mm.xformers_enabled.return_value = False + mock_mm.xformers_enabled_vae.return_value = False + mock_mm.sage_attention_enabled.return_value = False + mock_mm.flash_attention_enabled.return_value = False + + sentinel = object() + prior_cpu = cli_args.cpu + cli_args.cpu = True + prior_module = sys.modules.get("comfy_extras.nodes_seedvr", sentinel) + comfy_pkg = sys.modules.get("comfy") + prior_mm_attr = getattr(comfy_pkg, "model_management", sentinel) if comfy_pkg else sentinel + + with patch.dict(sys.modules, {"comfy.model_management": mock_mm}): + if comfy_pkg is not None: + setattr(comfy_pkg, "model_management", mock_mm) + sys.modules.pop("comfy_extras.nodes_seedvr", None) + try: + nodes_seedvr = importlib.import_module("comfy_extras.nodes_seedvr") + for node_cls in (nodes_seedvr.SeedVR2Preprocess, nodes_seedvr.SeedVR2PostProcessing, nodes_seedvr.SeedVR2Conditioning): + schema_ids = [i.id for i in node_cls.define_schema().inputs] + exec_params = [ + p for p in inspect.signature(node_cls.execute).parameters.keys() + if p != "cls" + ] + assert schema_ids == exec_params, ( + f"{node_cls.__name__} schema/execute drift: " + f"schema_ids={schema_ids}, exec_params={exec_params}" + ) + finally: + cli_args.cpu = prior_cpu + if prior_module is sentinel: + sys.modules.pop("comfy_extras.nodes_seedvr", None) + else: + sys.modules["comfy_extras.nodes_seedvr"] = prior_module + if comfy_pkg is not None: + if prior_mm_attr is sentinel: + if hasattr(comfy_pkg, "model_management"): + delattr(comfy_pkg, "model_management") + else: + setattr(comfy_pkg, "model_management", prior_mm_attr) diff --git a/tests-unit/comfy_extras_test/test_seedvr2_post_processing.py b/tests-unit/comfy_extras_test/test_seedvr2_post_processing.py new file mode 100644 index 000000000..6c821136d --- /dev/null +++ b/tests-unit/comfy_extras_test/test_seedvr2_post_processing.py @@ -0,0 +1,51 @@ +from unittest.mock import patch + +import pytest +import torch + +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + +from comfy_extras import nodes_seedvr # noqa: E402 + + +def _schema_ids(items): + return [item.id for item in items] + + +def test_seedvr2_post_processing_schema(): + schema = nodes_seedvr.SeedVR2PostProcessing.define_schema() + + assert _schema_ids(schema.inputs) == ["images", "original_resized_images", "color_correction_method"] + assert schema.inputs[2].options == ["lab", "wavelet", "adain", "none"] + assert schema.inputs[2].default == "lab" + assert schema.outputs[0].get_io_type() == "IMAGE" + + +def test_seedvr2_post_processing_oom_error_uses_color_correction_method(monkeypatch): + decoded = torch.full((1, 3, 4, 4), 0.25) + reference = torch.full((1, 3, 4, 4), 0.75) + + def _lab(content, style): + raise torch.cuda.OutOfMemoryError("CUDA out of memory") + + monkeypatch.setattr(nodes_seedvr.comfy.model_management, "vae_device", lambda: torch.device("cpu")) + monkeypatch.setattr(nodes_seedvr.comfy.model_management, "get_free_memory", lambda device: 1_000_000) + + with patch.object(nodes_seedvr, "lab_color_transfer", _lab): + with pytest.raises(RuntimeError) as excinfo: + nodes_seedvr.SeedVR2PostProcessing._color_transfer_chunked( + decoded, reference, torch.device("cpu"), "lab", + ) + assert "color_correction_method=lab" in str(excinfo.value) + assert " method=lab" not in str(excinfo.value) + + +def test_seedvr2_post_processing_unknown_color_correction_method_raises(): + decoded = torch.zeros(1, 2, 4, 4, 3) + original = torch.zeros(1, 2, 4, 4, 3) + with pytest.raises(ValueError) as excinfo: + nodes_seedvr.SeedVR2PostProcessing.execute(decoded, original, "bogus") + assert "color_correction_method" in str(excinfo.value) diff --git a/tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py b/tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py new file mode 100644 index 000000000..328355b49 --- /dev/null +++ b/tests-unit/comfy_extras_test/test_seedvr2_temporal_chunk.py @@ -0,0 +1,77 @@ +"""SeedVR2 temporal chunk/merge node regression tests.""" + +import pytest +import torch + +from comfy.cli_args import args as cli_args +from comfy.ldm.seedvr.constants import ( + BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE, + SEEDVR2_CHUNK_GIB_PER_MPX_FRAME, + SEEDVR2_CHUNK_RESERVED_GIB, + SEEDVR2_CHUNK_SIGMA_GIB, + SEEDVR2_CHUNK_SIGMA_K, + SEEDVR2_LATENT_CHANNELS, +) + +if not torch.cuda.is_available(): + cli_args.cpu = True + +import comfy.model_management # noqa: E402 +from comfy_extras.nodes_seedvr import SeedVR2TemporalChunk, SeedVR2TemporalMerge, _seedvr2_chunk_crossfade_weights # noqa: E402 + +def _latent(t_latent, h=8, w=8, b=1): + g = torch.Generator().manual_seed(7) + return {"samples": torch.randn(b, SEEDVR2_LATENT_CHANNELS, t_latent, h, w, generator=g)} + +def _split(latent, frames_per_chunk, temporal_overlap, chunking_mode="manual"): + combo = {"chunking_mode": chunking_mode} + if chunking_mode != "auto": + combo["frames_per_chunk"] = frames_per_chunk + return SeedVR2TemporalChunk.execute(latent, temporal_overlap, combo).args + +def _merge(chunks, temporal_overlap): + return SeedVR2TemporalMerge.execute(chunks, [temporal_overlap]).args[0] + +def test_chunk_temporal_windows_and_validation(): + with pytest.raises(ValueError, match="4n\\+1"): + _split(_latent(9), 20, 0) + with pytest.raises(ValueError, match="5-D"): + _split({"samples": torch.zeros(1, SEEDVR2_LATENT_CHANNELS * 9, 8, 8)}, 21, 0) + with pytest.raises(ValueError, match="chunking_mode"): + _split(_latent(13), 21, 0, "adaptive") + latent = _latent(13) + chunks, overlap = _split(latent, 21, 2) # chunk_latent=6, step=4 -> [0:6], [4:10], [8:13] + assert overlap == 2 and [c["samples"].shape[2] for c in chunks] == [6, 6, 5] + assert all(torch.equal(c["samples"], latent["samples"][:, :, s:e]) for c, (s, e) in zip(chunks, [(0, 6), (4, 10), (8, 13)])) + assert len(_split(_latent(13), 21, 999)[0]) == 8 # overlap clamps to chunk_latent-1 -> step=1 + assert (r := _split(_latent(5), 21, 3)) and len(r[0]) == 1 and r[1] == 0 # t_pixel <= 21: passthrough + +def test_chunk_auto_mode_applies_vram_law(monkeypatch): + mpx_per_frame = (32 * 32) * (BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE ** 2) / 1e6 + free_gb = ( + SEEDVR2_CHUNK_RESERVED_GIB + + SEEDVR2_CHUNK_SIGMA_K * SEEDVR2_CHUNK_SIGMA_GIB + + 5.1 * SEEDVR2_CHUNK_GIB_PER_MPX_FRAME * mpx_per_frame + ) + monkeypatch.setattr(comfy.model_management, "get_free_memory", lambda dev=None: free_gb * (1024 ** 3)) + assert [c["samples"].shape[2] for c in _split(_latent(13, h=32, w=32), 1, 0, "auto")[0]] == [5, 5, 3] + assert _split(_latent(13, h=32, w=32, b=2), 1, 0, "auto")[0][0]["samples"].shape[2] == 2 # batch halves the chunk + +def test_merge_crossfade_and_reassembly(): + latent = _latent(13) + latent["noise_mask"] = torch.rand(1, 1, 13, 8, 8) + latent["batch_index"] = [0] + merged = _merge(_split(latent, 21, 0)[0], 0) + assert torch.equal(merged["samples"], latent["samples"]) + assert "noise_mask" not in merged and merged["batch_index"] == [0] + assert torch.allclose(_merge(_split(latent, 21, 3)[0], 3)["samples"], latent["samples"], atol=1e-6) + w = _seedvr2_chunk_crossfade_weights(3, merged["samples"].device, merged["samples"].dtype) + assert w[0] == 1.0 and w[-1] == 0.0 and torch.all(w[:-1] >= w[1:]) + ones, zeros = {"samples": torch.ones(1, SEEDVR2_LATENT_CHANNELS, 6, 8, 8)}, {"samples": torch.zeros(1, SEEDVR2_LATENT_CHANNELS, 6, 8, 8)} + fused = _merge([ones, zeros], 3)["samples"] # overlap equals w: prev fades out, next fades in + assert torch.equal(fused[:, :, 3:6], w.view(1, 1, 3, 1, 1).expand(1, SEEDVR2_LATENT_CHANNELS, 3, 8, 8)) + assert torch.equal(fused[:, :, :3], ones["samples"][:, :, :3]) and torch.equal(fused[:, :, 6:], zeros["samples"][:, :, :3]) + short = _split(latent, 21, 2)[0] + short[0]["samples"] = short[0]["samples"][:, :, :4] + with pytest.raises(ValueError, match="only the final chunk may be shorter"): + _merge(short, 2) diff --git a/tests-unit/comfy_test/model_detection_test.py b/tests-unit/comfy_test/model_detection_test.py index 4e9350602..6e7d71f79 100644 --- a/tests-unit/comfy_test/model_detection_test.py +++ b/tests-unit/comfy_test/model_detection_test.py @@ -2,7 +2,7 @@ from collections import defaultdict import torch -from comfy.model_detection import detect_unet_config, model_config_from_unet_config +from comfy.model_detection import detect_unet_config, model_config_from_unet, model_config_from_unet_config import comfy.supported_models @@ -73,6 +73,34 @@ def _make_flux_schnell_comfyui_sd(): return sd +def _make_seedvr2_7b_separate_mm_sd(): + return { + "blocks.35.mlp.vid.proj_out.weight": torch.empty(3072, 1), + "positive_conditioning": torch.empty(58, 5120), + "negative_conditioning": torch.empty(64, 5120), + } + + +def _make_seedvr2_7b_shared_mm_sd(): + return { + "blocks.35.mlp.all.proj_in_gate.weight": torch.empty(1, 1), + "positive_conditioning": torch.empty(58, 5120), + "negative_conditioning": torch.empty(64, 5120), + } + + +def _make_seedvr2_3b_shared_mm_sd(): + return { + "blocks.31.mlp.all.proj_in_gate.weight": torch.empty(1, 1), + "positive_conditioning": torch.empty(58, 5120), + "negative_conditioning": torch.empty(64, 5120), + } + + +def _add_model_diffusion_prefix(sd): + return {f"model.diffusion_model.{k}": v for k, v in sd.items()} + + class TestModelDetection: """Verify that first-match model detection selects the correct model based on list ordering and unet_config specificity.""" @@ -125,6 +153,59 @@ class TestModelDetection: assert model_config is not None assert type(model_config).__name__ == "FluxSchnell" + def test_seedvr2_7b_separate_mm_detection_config(self): + sd = _make_seedvr2_7b_separate_mm_sd() + unet_config = detect_unet_config(sd, "") + + assert unet_config is not None + assert unet_config["image_model"] == "seedvr2" + assert unet_config["vid_dim"] == 3072 + assert unet_config["heads"] == 24 + assert unet_config["num_layers"] == 36 + assert unet_config["mm_layers"] == 36 + assert unet_config["mlp_type"] == "normal" + assert unet_config["rope_type"] == "rope3d" + assert unet_config["rope_dim"] == 64 + + def test_seedvr2_7b_shared_mm_detection_config(self): + sd = _make_seedvr2_7b_shared_mm_sd() + unet_config = detect_unet_config(sd, "") + + assert unet_config is not None + assert unet_config["image_model"] == "seedvr2" + assert unet_config["vid_dim"] == 3072 + assert unet_config["heads"] == 24 + assert unet_config["num_layers"] == 36 + assert unet_config["mm_layers"] == 10 + assert unet_config["mlp_type"] == "swiglu" + assert unet_config["rope_type"] == "rope3d" + assert unet_config["rope_dim"] == 64 + + def test_seedvr2_3b_shared_mm_detection_config(self): + sd = _make_seedvr2_3b_shared_mm_sd() + unet_config = detect_unet_config(sd, "") + + assert unet_config is not None + assert unet_config["image_model"] == "seedvr2" + assert unet_config["vid_dim"] == 2560 + assert unet_config["heads"] == 20 + assert unet_config["num_layers"] == 32 + assert unet_config["mlp_type"] == "swiglu" + + def test_seedvr2_model_match_requires_conditioning_tensors(self): + sd = _make_seedvr2_7b_shared_mm_sd() + unet_config = detect_unet_config(sd, "") + + assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "SeedVR2" + + del sd["positive_conditioning"] + assert model_config_from_unet_config(unet_config, sd) is None + + def test_seedvr2_model_match_accepts_full_checkpoint_prefix(self): + sd = _add_model_diffusion_prefix(_make_seedvr2_7b_shared_mm_sd()) + + assert type(model_config_from_unet(sd, "model.diffusion_model.")).__name__ == "SeedVR2" + def test_unet_config_and_required_keys_combination_is_unique(self): """Each model in the registry must have a unique combination of ``unet_config`` and ``required_keys``. If two models share the same diff --git a/tests-unit/comfy_test/seedvr_vae_forward_test.py b/tests-unit/comfy_test/seedvr_vae_forward_test.py new file mode 100644 index 000000000..7ea7a143e --- /dev/null +++ b/tests-unit/comfy_test/seedvr_vae_forward_test.py @@ -0,0 +1,74 @@ +"""Regression tests for the SeedVR2 VAE forward return contract.""" + +import pytest +import torch +import torch.nn as nn + +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + +from comfy.ldm.seedvr.vae import SEEDVR2_LATENT_CHANNELS, VideoAutoencoderKL # noqa: E402 + + +_LATENT_SHAPE = (1, SEEDVR2_LATENT_CHANNELS, 2, 2, 2) +_DECODED_SHAPE = (1, 3, 5, 16, 16) +_INPUT_ENCODE_SHAPE = (1, 3, 5, 16, 16) +_INPUT_DECODE_SHAPE = _LATENT_SHAPE + + +class _StubVAE(VideoAutoencoderKL): + def __init__(self): + nn.Module.__init__(self) + self._encode_out = torch.zeros(*_LATENT_SHAPE) + self._decode_out = torch.zeros(*_DECODED_SHAPE) + + def encode(self, x, return_dict=True): + return self._encode_out + + def decode_(self, z, return_dict=True): + return self._decode_out + + +def test_forward_encode_returns_tensor(): + vae = _StubVAE() + x = torch.zeros(*_INPUT_ENCODE_SHAPE) + result = vae.forward(x, mode="encode") + assert type(result) is torch.Tensor + assert result.shape == torch.Size(_LATENT_SHAPE) + + +def test_forward_decode_returns_tensor(): + vae = _StubVAE() + z = torch.zeros(*_INPUT_DECODE_SHAPE) + result = vae.forward(z, mode="decode") + assert type(result) is torch.Tensor + assert result.shape == torch.Size(_DECODED_SHAPE) + + +class _TupleReturningStubVAE(VideoAutoencoderKL): + def __init__(self): + nn.Module.__init__(self) + self._encode_tensor = torch.zeros(*_LATENT_SHAPE) + self._decode_tensor = torch.zeros(*_DECODED_SHAPE) + + def encode(self, x, return_dict=True): + return (self._encode_tensor,) + + def decode_(self, z, return_dict=True): + return (self._decode_tensor,) + + +def test_forward_all_unwraps_one_tuple_at_each_step(): + vae = _TupleReturningStubVAE() + x = torch.zeros(*_INPUT_ENCODE_SHAPE) + result = vae.forward(x, mode="all") + assert type(result) is torch.Tensor + assert result.shape == torch.Size(_DECODED_SHAPE) + + +def test_forward_rejects_unknown_mode(): + vae = _StubVAE() + with pytest.raises(ValueError, match="Unknown SeedVR2 VAE forward mode"): + vae.forward(torch.zeros(*_INPUT_ENCODE_SHAPE), mode="bogus") diff --git a/tests-unit/comfy_test/test_seedvr2_dtype.py b/tests-unit/comfy_test/test_seedvr2_dtype.py new file mode 100644 index 000000000..8e08b6dde --- /dev/null +++ b/tests-unit/comfy_test/test_seedvr2_dtype.py @@ -0,0 +1,50 @@ +import torch + +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + +import comfy.sd +import comfy.supported_models +import comfy.ldm.seedvr.model as seedvr_model +import comfy.ldm.seedvr.vae as seedvr_vae + + +def test_seedvr2_fp16_manual_cast_only_for_bf16_device(monkeypatch): + bf16_device = object() + fp16_device = object() + + monkeypatch.setattr( + comfy.supported_models.comfy.model_management, + "should_use_bf16", + lambda device=None: device is bf16_device, + ) + + bf16_config = comfy.supported_models.SeedVR2({"image_model": "seedvr2"}) + bf16_config.set_inference_dtype(torch.float16, None, device=bf16_device) + assert bf16_config.manual_cast_dtype is torch.bfloat16 + + fp16_config = comfy.supported_models.SeedVR2({"image_model": "seedvr2"}) + fp16_config.set_inference_dtype(torch.float16, None, device=fp16_device) + assert fp16_config.manual_cast_dtype is None + + +def test_seedvr2_text_conditioning_accepts_cfg1_single_branch(): + context = torch.arange(6, dtype=torch.float32).reshape(1, 3, 2) + + txt, txt_shape = seedvr_model.NaDiT._resolve_text_conditioning(object(), context, [0]) + + torch.testing.assert_close(txt, context.squeeze(0)) + torch.testing.assert_close(txt_shape, torch.tensor([[3]], device=context.device)) + + +def test_seedvr2_vae_decode_memory_covers_full_frame_lab_transfer(): + wrapper = seedvr_vae.VideoAutoencoderKLWrapper.__new__(seedvr_vae.VideoAutoencoderKLWrapper) + latent_channels = seedvr_vae.SEEDVR2_LATENT_CHANNELS + estimate = wrapper.comfy_memory_used_decode((1, latent_channels, 26, 120, 160)) + old_estimate = latent_channels * 120 * 160 * (4 * 8 * 8) * 2 + + assert estimate == 101 * 960 * 1280 * 160 + assert estimate > 15 * 1024 ** 3 + assert estimate > old_estimate * 100 diff --git a/tests-unit/comfy_test/test_seedvr2_internals.py b/tests-unit/comfy_test/test_seedvr2_internals.py new file mode 100644 index 000000000..fe4bde1c4 --- /dev/null +++ b/tests-unit/comfy_test/test_seedvr2_internals.py @@ -0,0 +1,169 @@ +"""SeedVR2 internals regression tests.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +import torch + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +import comfy.ldm.seedvr.model as seedvr_model # noqa: E402 +import comfy.ldm.seedvr.vae as vae_mod # noqa: E402 +import comfy.ldm.modules.attention as attention # noqa: E402 +import comfy.ops as comfy_ops # noqa: E402 +from comfy.ldm.seedvr.vae import ( # noqa: E402 + causal_norm_wrapper, + set_norm_limit, +) +from comfy.ldm.seedvr.attention import var_attention_optimized_split # noqa: E402 + + +_NUM_CHANNELS = 8 +_NUM_GROUPS = 4 +_TENSOR_SHAPE = (1, 8, 2, 4, 4) + +_GROUPNORM_SUBCLASSES = [ + pytest.param(comfy_ops.disable_weight_init.GroupNorm, id="disable_weight_init"), + pytest.param(comfy_ops.manual_cast.GroupNorm, id="manual_cast"), +] + + +@pytest.mark.parametrize("groupnorm_cls", _GROUPNORM_SUBCLASSES) +def test_seedvr_groupnorm_low_limit_uses_chunked_groupnorm_path(groupnorm_cls): + real_group_norm = vae_mod.F.group_norm + set_norm_limit(1e-9) + try: + gn = groupnorm_cls(num_channels=_NUM_CHANNELS, num_groups=_NUM_GROUPS) + gn.eval() + + forward_hook_calls = [] + + def _hook(module, inputs, output): + forward_hook_calls.append(tuple(inputs[0].shape)) + + spy_calls = [] + + def _group_norm_spy(input_tensor, num_groups_arg, *args, **kwargs): + spy_calls.append({"num_groups": int(num_groups_arg)}) + return real_group_norm(input_tensor, num_groups_arg, *args, **kwargs) + + handle = gn.register_forward_hook(_hook) + try: + with patch.object(vae_mod.F, "group_norm", side_effect=_group_norm_spy): + out_tensor = causal_norm_wrapper(gn, torch.randn(*_TENSOR_SHAPE)) + finally: + handle.remove() + + full_calls = len(forward_hook_calls) + chunked_calls = sum(1 for entry in spy_calls if entry["num_groups"] < _NUM_GROUPS) + + assert tuple(int(s) for s in out_tensor.shape) == _TENSOR_SHAPE + assert full_calls == 0, ( + f"low-limit GroupNorm gate must NOT take the full-forward path; got full_calls={full_calls}" + ) + assert chunked_calls > 0, ( + f"low-limit GroupNorm gate must take the chunked path; got chunked_calls={chunked_calls}" + ) + finally: + set_norm_limit(None) + + +def test_seedvr2_7b_swin_attention_forward_uses_optimized_var_attention(monkeypatch): + dim = 8 + heads = 2 + head_dim = 4 + attn = seedvr_model.NaSwinAttention( + vid_dim=dim, + txt_dim=dim, + heads=heads, + head_dim=head_dim, + qk_bias=False, + qk_norm=comfy_ops.disable_weight_init.RMSNorm, + qk_norm_eps=1e-6, + rope_type=None, + rope_dim=head_dim, + shared_weights=False, + window=(2, 1, 1), + window_method="720pwin_by_size_bysize", + version=True, + device="cpu", + dtype=torch.float32, + operations=comfy_ops.disable_weight_init, + ) + generator = torch.Generator(device="cpu").manual_seed(11) + vid = torch.randn(8, dim, generator=generator) + txt = torch.randn(3, dim, generator=generator) + vid_shape = torch.tensor([[2, 2, 2]], dtype=torch.long) + txt_shape = torch.tensor([[3]], dtype=torch.long) + calls = [] + + def fake_optimized_var_attention(**kwargs): + calls.append(kwargs) + return kwargs["q"] + + monkeypatch.setattr(seedvr_model, "optimized_var_attention", fake_optimized_var_attention) + + vid_out, txt_out = attn(vid, txt, vid_shape, txt_shape, seedvr_model.Cache(disable=True)) + + assert tuple(vid_out.shape) == (8, dim) + assert tuple(txt_out.shape) == (3, dim) + assert len(calls) == 1 + call = calls[0] + assert tuple(call["q"].shape) == (14, heads, head_dim) + assert tuple(call["k"].shape) == (14, heads, head_dim) + assert tuple(call["v"].shape) == (14, heads, head_dim) + assert call["heads"] == heads + assert call["skip_reshape"] is True + assert call["skip_output_reshape"] is True + assert call["cu_seqlens_q"] == [0, 7, 14] + assert call["cu_seqlens_k"] == [0, 7, 14] + + +def test_var_attention_optimized_split_calls_dense_backend_per_window(monkeypatch): + heads = 2 + head_dim = 3 + q = torch.arange(30, dtype=torch.float32).reshape(5, heads, head_dim) + k = q + 100 + v = q + 200 + cu = [0, 2, 5] + calls = [] + + def fake_optimized_attention(q_arg, k_arg, v_arg, heads_arg, **kwargs): + calls.append( + { + "q_shape": tuple(q_arg.shape), + "k_shape": tuple(k_arg.shape), + "v_shape": tuple(v_arg.shape), + "heads": heads_arg, + "kwargs": kwargs, + } + ) + return q_arg + v_arg + + monkeypatch.setattr(attention, "optimized_attention", fake_optimized_attention) + + out = var_attention_optimized_split( + q, + k, + v, + heads, + cu, + cu, + skip_reshape=True, + skip_output_reshape=True, + ) + + assert tuple(out.shape) == (5, heads, head_dim) + assert len(calls) == 2 + assert calls[0]["q_shape"] == (1, heads, 2, head_dim) + assert calls[1]["q_shape"] == (1, heads, 3, head_dim) + assert all(call["heads"] == heads for call in calls) + assert all(call["kwargs"]["skip_reshape"] is True for call in calls) + assert all(call["kwargs"]["skip_output_reshape"] is True for call in calls) + torch.testing.assert_close(out, q + v, rtol=0, atol=0) + diff --git a/tests-unit/comfy_test/test_seedvr2_model.py b/tests-unit/comfy_test/test_seedvr2_model.py new file mode 100644 index 000000000..1d454aaf1 --- /dev/null +++ b/tests-unit/comfy_test/test_seedvr2_model.py @@ -0,0 +1,320 @@ +"""SeedVR2 model, latent-format, and VAE graph regression tests.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest +import torch +from torch import nn + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +import comfy # noqa: E402 +import comfy.latent_formats # noqa: E402 +import comfy.ldm.seedvr.model as seedvr_model # noqa: E402 +import comfy.ldm.seedvr.vae as seedvr_vae_mod # noqa: E402 +import comfy.model_management # noqa: E402 +import comfy.ops as comfy_ops # noqa: E402 +import comfy.sample # noqa: E402 +import comfy.sd as sd_mod # noqa: E402 +import nodes as nodes_mod # noqa: E402 +from comfy.ldm.seedvr.model import NaDiT # noqa: E402 + + +_LATENT_CHANNELS = seedvr_vae_mod.SEEDVR2_LATENT_CHANNELS + + +def _make_standin(positive_conditioning): + class _StandIn(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "positive_conditioning", positive_conditioning + ) + + _resolve_text_conditioning = NaDiT._resolve_text_conditioning + + return _StandIn() + + +class _StubModule(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + +def _capture_last_layer_flags(monkeypatch, vid_dim: int, txt_in_dim: int) -> list[bool]: + flags = [] + + class _Block(_StubModule): + def __init__(self, *args, **kwargs): + flags.append(kwargs["is_last_layer"]) + super().__init__() + + monkeypatch.setattr(seedvr_model, "NaPatchIn", _StubModule) + monkeypatch.setattr(seedvr_model, "NaPatchOut", _StubModule) + monkeypatch.setattr(seedvr_model, "TimeEmbedding", _StubModule) + monkeypatch.setattr(seedvr_model, "NaMMSRTransformerBlock", _Block) + + seedvr_model.NaDiT( + norm_eps=1e-5, + num_layers=4, + mlp_type="normal", + vid_dim=vid_dim, + txt_in_dim=txt_in_dim, + heads=24, + mm_layers=3, + operations=comfy_ops.disable_weight_init, + ) + + return flags + + +class _Model: + def __init__(self, latent_format): + self._latent_format = latent_format + + def get_model_object(self, name): + assert name == "latent_format" + return self._latent_format + + +class _Patcher: + def get_free_memory(self, device): + return 1024 * 1024 * 1024 + + +class _EncodeWrapper(seedvr_vae_mod.VideoAutoencoderKLWrapper): + def __init__(self, encoded): + nn.Module.__init__(self) + self.encoded = encoded + self.spatial_downsample_factor = 8 + self.temporal_downsample_factor = 4 + self.seen = [] + + def encode(self, x): + self.seen.append(tuple(x.shape)) + return self.encoded.to(device=x.device, dtype=x.dtype) + + +class _DecodeWrapper(seedvr_vae_mod.VideoAutoencoderKLWrapper): + def __init__(self): + nn.Module.__init__(self) + self.spatial_downsample_factor = 8 + self.temporal_downsample_factor = 4 + self.calls = [] + + def decode(self, z, seedvr2_tiling=None): + self.calls.append({"shape": tuple(z.shape), "seedvr2_tiling": seedvr2_tiling}) + if z.ndim == 4: + b, tc, h, w = z.shape + t = tc // _LATENT_CHANNELS + else: + b, _, t, h, w = z.shape + return torch.zeros(b, 3, t, h * 8, w * 8, dtype=z.dtype, device=z.device) + + +def test_seedvr2_wrapper_public_encode_returns_tensor(monkeypatch): + raw_latent = torch.full((1, _LATENT_CHANNELS, 1, 4, 5), 2.0) + seen_shapes = [] + + def base_encode(self, x): + seen_shapes.append(tuple(x.shape)) + return raw_latent.to(device=x.device, dtype=x.dtype) + + monkeypatch.setattr(seedvr_vae_mod.VideoAutoencoderKL, "encode", base_encode) + + vae = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__(seedvr_vae_mod.VideoAutoencoderKLWrapper) + nn.Module.__init__(vae) + vae._dummy = nn.Parameter(torch.zeros((), dtype=torch.float32)) + + latent = vae.encode(torch.zeros(1, 3, 32, 40)) + + assert type(latent) is torch.Tensor + assert tuple(latent.shape) == (1, _LATENT_CHANNELS, 4, 5) + assert seen_shapes == [(1, 3, 1, 32, 40)] + + +def test_seedvr2_wrapper_private_encode_helper_keeps_raw_latent(monkeypatch): + raw_latent = torch.full((1, _LATENT_CHANNELS, 1, 4, 5), 3.0) + + def base_encode(self, x): + return raw_latent.to(device=x.device, dtype=x.dtype) + + monkeypatch.setattr(seedvr_vae_mod.VideoAutoencoderKL, "encode", base_encode) + + vae = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__(seedvr_vae_mod.VideoAutoencoderKLWrapper) + nn.Module.__init__(vae) + vae._dummy = nn.Parameter(torch.zeros((), dtype=torch.float32)) + + latent, raw = vae._encode_with_raw_latent(torch.zeros(1, 3, 32, 40)) + + assert tuple(latent.shape) == (1, _LATENT_CHANNELS, 4, 5) + assert tuple(raw.shape) == (1, _LATENT_CHANNELS, 1, 4, 5) + assert torch.equal(raw, raw_latent) + + +def _make_vae(wrapper): + vae = sd_mod.VAE.__new__(sd_mod.VAE) + vae.first_stage_model = wrapper + vae.device = torch.device("cpu") + vae.output_device = torch.device("cpu") + vae.vae_dtype = torch.float32 + vae.latent_channels = _LATENT_CHANNELS + vae.latent_dim = 3 + vae.downscale_ratio = (lambda a: max(0, (a + 3) // 4), 8, 8) + vae.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8) + vae.output_channels = 3 + vae.disable_offload = True + vae.extra_1d_channel = None + vae.crop_input = False + vae.not_video = False + vae.handles_tiling = isinstance(wrapper, seedvr_vae_mod.VideoAutoencoderKLWrapper) + vae.format_encoded = wrapper.comfy_format_encoded + vae.patcher = _Patcher() + vae.process_input = lambda image: image + vae.process_output = lambda image: image.add(1.0).div(2.0).clamp(0.0, 1.0) + vae.vae_output_dtype = lambda: torch.float32 + vae.memory_used_encode = lambda shape, dtype: 1 + vae.memory_used_decode = lambda shape, dtype: 1 + vae.throw_exception_if_invalid = lambda: None + vae.vae_encode_crop_pixels = lambda pixels: pixels + vae.spacial_compression_decode = lambda: 8 + vae.temporal_compression_decode = lambda: 4 + return vae + + +def test_missing_context_falls_back_to_positive_buffer(): + pos_buffer = torch.full((58, 5120), 7.0) + standin = _make_standin(pos_buffer) + txt, txt_shape = standin._resolve_text_conditioning(None) + assert txt.shape == (58, 5120) + assert (txt == 7.0).all(), ( + "fallback path must use the positive_conditioning buffer " + "verbatim, not a zero tensor" + ) + assert txt_shape.shape == (1, 1) + assert txt_shape[0, 0].item() == 58 + + +def test_seedvr2_7b_keeps_final_block_text_path(monkeypatch): + assert _capture_last_layer_flags(monkeypatch, vid_dim=3072, txt_in_dim=3072) == [ + False, + False, + False, + False, + ] + + +def test_seedvr2_7b_rope3d_matches_wrapper_oracle(): + rope = seedvr_model.get_na_rope("rope3d", dim=64) + generator = torch.Generator(device="cpu").manual_seed(0) + q = torch.randn(4, 2, 128, generator=generator) + k = torch.randn(4, 2, 128, generator=generator) + shape = torch.tensor([[1, 2, 2]], dtype=torch.long) + freqs = rope.get_axial_freqs(1, 2, 2).reshape(4, -1) + + expected_q = seedvr_model._apply_seedvr2_rotary_emb( + freqs, + q.permute(1, 0, 2).float(), + ).to(q.dtype).permute(1, 0, 2) + expected_k = seedvr_model._apply_seedvr2_rotary_emb( + freqs, + k.permute(1, 0, 2).float(), + ).to(k.dtype).permute(1, 0, 2) + + actual_q, actual_k = rope(q.clone(), k.clone(), shape, seedvr_model.Cache(disable=True)) + + torch.testing.assert_close(actual_q, expected_q, rtol=0, atol=0) + torch.testing.assert_close(actual_k, expected_k, rtol=0, atol=0) + + +def test_seedvr2_forward_requires_conditioning_latents(): + model = NaDiT.__new__(NaDiT) + x = torch.zeros(1, _LATENT_CHANNELS, 1, 4, 5) + + with pytest.raises(ValueError, match="requires conditioning latents"): + NaDiT.forward(model, x, timestep=torch.tensor([1.0]), context=None) + + +def test_seedvr2_latent_format_uses_native_video_latent_shape(): + latent_format = comfy.latent_formats.SeedVR2() + latent_image = torch.zeros(1, 1, 4, 5) + + fixed = comfy.sample.fix_empty_latent_channels(_Model(latent_format), latent_image) + + assert latent_format.latent_channels == _LATENT_CHANNELS + assert latent_format.latent_dimensions == 3 + assert fixed.shape == (1, _LATENT_CHANNELS, 1, 4, 5) + + +def test_seedvr2_model_requires_native_5d_latent(): + latent = torch.zeros(1, _LATENT_CHANNELS, 2, 4, 5) + assert NaDiT._check_seedvr2_video_latent(latent, _LATENT_CHANNELS, "latent") is latent + + with pytest.raises(ValueError, match="5-D native latent"): + NaDiT._check_seedvr2_video_latent(torch.zeros(1, _LATENT_CHANNELS * 2, 4, 5), _LATENT_CHANNELS, "latent") + + +def test_seedvr2_encode_and_encode_tiled_preserve_native_latent_contract(monkeypatch): + monkeypatch.setattr(sd_mod.model_management, "load_models_gpu", lambda *a, **k: None) + + encoded = torch.full((1, _LATENT_CHANNELS, 2, 4, 5), 2.0) + vae = _make_vae(_EncodeWrapper(encoded)) + pixels = torch.zeros(1, 5, 32, 40, 3) + + node_output = nodes_mod.VAEEncode().encode(vae, pixels)[0] + node_latent = node_output["samples"] + assert set(node_output) == {"samples"} + assert tuple(node_latent.shape) == (1, _LATENT_CHANNELS, 2, 4, 5) + assert node_latent.dtype == torch.float32 + assert node_latent.stride()[-1] == 1 + assert torch.equal(node_latent, torch.full_like(node_latent, 2.0 * seedvr_vae_mod.BYTEDANCE_VAE_SCALING_FACTOR)) + + tiled = torch.full((1, _LATENT_CHANNELS, 2, 4, 5), 3.0) + monkeypatch.setattr(seedvr_vae_mod, "tiled_vae", MagicMock(return_value=tiled)) + tiled_output = nodes_mod.VAEEncodeTiled().encode( + vae, + pixels, + tile_size=512, + overlap=64, + temporal_size=16, + temporal_overlap=4, + )[0] + tiled_latent = tiled_output["samples"] + assert set(tiled_output) == {"samples"} + assert tuple(tiled_latent.shape) == (1, _LATENT_CHANNELS, 2, 4, 5) + assert tiled_latent.dtype == torch.float32 + assert torch.equal(tiled_latent, torch.full_like(tiled_latent, 3.0 * seedvr_vae_mod.BYTEDANCE_VAE_SCALING_FACTOR)) + + +def test_vaedecode_tiled_spatial_applies_temporal_discarded(monkeypatch): + monkeypatch.setattr(sd_mod.model_management, "load_models_gpu", lambda *a, **k: None) + vae = _make_vae(_DecodeWrapper()) + + nodes_mod.VAEDecodeTiled().decode( + vae, + {"samples": torch.zeros(1, _LATENT_CHANNELS, 2, 4, 5)}, + tile_size=512, + overlap=64, + temporal_size=16, + temporal_overlap=4, + ) + + # Spatial inputs flow through; temporal inputs are discarded as public tiling + # knobs, but SeedVR2's internal MemoryState causal slicing is left intact. + assert vae.first_stage_model.calls == [ + { + "shape": (1, _LATENT_CHANNELS, 2, 4, 5), + "seedvr2_tiling": { + "enable_tiling": True, + "tile_size": (512, 512), + "tile_overlap": (64, 64), + "temporal_size": None, + "temporal_overlap": None, + }, + } + ] diff --git a/tests-unit/comfy_test/test_seedvr2_vae_decode.py b/tests-unit/comfy_test/test_seedvr2_vae_decode.py new file mode 100644 index 000000000..c486b9195 --- /dev/null +++ b/tests-unit/comfy_test/test_seedvr2_vae_decode.py @@ -0,0 +1,94 @@ +from unittest.mock import patch + +import pytest +import torch +import torch.nn as nn + +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + +import comfy.ldm.seedvr.vae as vae_mod # noqa: E402 +from comfy_extras import nodes_seedvr # noqa: E402 + + +_LATENT_CHANNELS = vae_mod.SEEDVR2_LATENT_CHANNELS + + +def _make_wrapper() -> vae_mod.VideoAutoencoderKLWrapper: + wrapper = vae_mod.VideoAutoencoderKLWrapper.__new__( + vae_mod.VideoAutoencoderKLWrapper + ) + nn.Module.__init__(wrapper) + return wrapper + + +def _fingerprint_decode_(self, z, return_dict=True): + b = int(z.shape[0]) + t = int(z.shape[2]) + h = int(z.shape[3]) + w = int(z.shape[4]) + out = torch.empty(b, 3, t, h * 8, w * 8) + for batch_idx in range(b): + out[batch_idx].fill_(float(batch_idx + 1)) + return out + + +def _decode_with_patches(wrapper, z): + with patch.object(vae_mod.VideoAutoencoderKL, "decode_", _fingerprint_decode_): + return wrapper.decode(z) + + +def test_decode_b2_t3_multi_frame_batch_unchanged(): + wrapper = _make_wrapper() + + out = _decode_with_patches(wrapper, torch.zeros(2, _LATENT_CHANNELS * 3, 2, 2)) + + assert tuple(out.shape) == (2, 3, 3, 16, 16) + + +class _Wrapper(vae_mod.VideoAutoencoderKLWrapper): + def __init__(self): + nn.Module.__init__(self) + self.calls = [] + + def parameters(self): + return iter([torch.nn.Parameter(torch.zeros(()))]) + +def _decode_stub(self, latent): + self.calls.append(tuple(latent.shape)) + return torch.zeros(latent.shape[0], 3, latent.shape[2], latent.shape[3] * 8, latent.shape[4] * 8) + + +def test_seedvr2_wrapper_decode_accepts_5d_channel_first_latents_without_preprocessor_state(): + wrapper = _Wrapper() + + with patch.object(vae_mod.VideoAutoencoderKL, "decode_", _decode_stub): + out = wrapper.decode(torch.zeros(1, _LATENT_CHANNELS, 2, 4, 5)) + + assert tuple(out.shape) == (1, 3, 2, 32, 40) + assert wrapper.calls == [(1, _LATENT_CHANNELS, 2, 4, 5)] + + +def test_seedvr2_wrapper_decode_rejects_wrong_rank_latents(): + wrapper = _Wrapper() + + with pytest.raises(RuntimeError, match=r"latent input must be 4-D collapsed .* or 5-D"): + wrapper.decode(torch.zeros(1, _LATENT_CHANNELS, 4)) + + +def _t_padded(t_in: int) -> int: + if t_in == 1: + return 1 + if t_in <= 4: + return 5 + if (t_in - 1) % 4 == 0: + return t_in + return t_in + (4 - ((t_in - 1) % 4)) + + +@pytest.mark.parametrize("t_in", [1, 5, 9]) +def test_t_padded_matches_cut_videos(t_in): + dummy = torch.zeros(1, t_in, 1, 1, 1) + assert nodes_seedvr.cut_videos(dummy).shape[1] == _t_padded(t_in) diff --git a/tests-unit/comfy_test/test_seedvr2_vae_tiled.py b/tests-unit/comfy_test/test_seedvr2_vae_tiled.py new file mode 100644 index 000000000..a2866b609 --- /dev/null +++ b/tests-unit/comfy_test/test_seedvr2_vae_tiled.py @@ -0,0 +1,382 @@ +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.nn as nn + +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + +import comfy.ldm.seedvr.vae as vae_mod # noqa: E402 +import comfy.ldm.seedvr.vae as seedvr_vae_mod # noqa: E402 +import comfy.sd as sd_mod # noqa: E402 +from comfy.ldm.seedvr.vae import MemoryState, tiled_vae # noqa: E402 + + +_LATENT_CHANNELS = seedvr_vae_mod.SEEDVR2_LATENT_CHANNELS + + +def test_runtime_decode_zero_temporal_size_preserves_model_slicing(): + class StubVAEModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.slicing_latent_min_size = 2 + self.spatial_downsample_factor = 8 + self.temporal_downsample_factor = 4 + self.device = torch.device("cpu") + self.use_slicing = True + self._dummy = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + self.decode_min_sizes = [] + self.memory_states = [] + + def decode_(self, t_chunk): + self.decode_min_sizes.append(self.slicing_latent_min_size) + return vae_mod.VideoAutoencoderKL.slicing_decode(self, t_chunk) + + def _decode(self, z, memory_state=MemoryState.DISABLED, memory_cache=None): + self.memory_states.append(memory_state) + b, c, d, h, w = z.shape + return torch.zeros((b, 3, d, h * 8, w * 8), dtype=z.dtype) + + vae = StubVAEModel() + z = torch.zeros((1, _LATENT_CHANNELS, 5, 8, 8), dtype=torch.float32) + + tiled_vae( + z, + vae, + tile_size=(64, 64), + tile_overlap=(0, 0), + temporal_size=0, + temporal_overlap=0, + encode=False, + ) + + assert vae.decode_min_sizes == [2] + assert vae.memory_states == [MemoryState.INITIALIZING, MemoryState.ACTIVE] + assert vae.slicing_latent_min_size == 2 + + +def test_zero_temporal_size_preserves_min_size_when_encode_raises(): + class RaisingVAEModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.slicing_sample_min_size = 4 + self.spatial_downsample_factor = 8 + self.temporal_downsample_factor = 4 + self.device = torch.device("cpu") + self._dummy = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + + def encode(self, t_chunk): + raise RuntimeError("simulated encode failure") + + vae = RaisingVAEModel() + x = torch.zeros((1, 3, 12, 64, 64), dtype=torch.float32) + + with pytest.raises(RuntimeError, match="simulated encode failure"): + tiled_vae( + x, + vae, + tile_size=(64, 64), + tile_overlap=(0, 0), + temporal_size=0, + temporal_overlap=0, + encode=True, + ) + + assert vae.slicing_sample_min_size == 4 + + +def test_tiled_vae_encode_uses_tensor_return_without_indexing(): + class TensorEncodeVAEModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.slicing_sample_min_size = 4 + self.spatial_downsample_factor = 8 + self.temporal_downsample_factor = 4 + self.device = torch.device("cpu") + self._dummy = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + self.calls = [] + + def encode(self, t_chunk): + self.calls.append(tuple(t_chunk.shape)) + b, _, _, h, w = t_chunk.shape + return torch.ones((b, _LATENT_CHANNELS, 1, h // 8, w // 8), dtype=t_chunk.dtype) + + vae = TensorEncodeVAEModel() + x = torch.zeros((2, 3, 1, 64, 64), dtype=torch.float32) + + out = tiled_vae( + x, + vae, + tile_size=(64, 64), + tile_overlap=(0, 0), + temporal_size=0, + temporal_overlap=0, + encode=True, + ) + + assert vae.calls == [(2, 3, 1, 64, 64)] + assert tuple(out.shape) == (2, _LATENT_CHANNELS, 1, 8, 8) + + +def test_tiled_vae_preserves_input_dtype_on_single_tile(): + class FloatOutputVAEModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.slicing_sample_min_size = 4 + self.spatial_downsample_factor = 8 + self.temporal_downsample_factor = 4 + self.device = torch.device("cpu") + self._dummy = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + + def encode(self, t_chunk): + b, _, _, h, w = t_chunk.shape + return torch.ones((b, _LATENT_CHANNELS, 1, h // 8, w // 8), dtype=torch.float32) + + out = tiled_vae( + torch.zeros((1, 3, 1, 64, 64), dtype=torch.float16), + FloatOutputVAEModel(), + tile_size=(64, 64), + tile_overlap=(0, 0), + temporal_size=0, + temporal_overlap=0, + encode=True, + ) + + assert out.dtype == torch.float16 + + +class _SlicingDecodeVAE(nn.Module): + def __init__(self, slicing_latent_min_size): + super().__init__() + self.slicing_latent_min_size = slicing_latent_min_size + self.spatial_downsample_factor = 8 + self.temporal_downsample_factor = 4 + self.device = torch.device("cpu") + self.use_slicing = True + self._dummy = nn.Parameter(torch.zeros(1, dtype=torch.float32)) + self.decode_min_sizes = [] + self.memory_states = [] + + def decode_(self, z): + self.decode_min_sizes.append(self.slicing_latent_min_size) + return vae_mod.VideoAutoencoderKL.slicing_decode(self, z) + + def _decode(self, z, memory_state=MemoryState.DISABLED, memory_cache=None): + self.memory_states.append(memory_state) + x = z[:, :1].repeat( + 1, + 3, + 1, + self.spatial_downsample_factor, + self.spatial_downsample_factor, + ) + return x + + +def test_decode_tiled_vae_maps_temporal_args_to_latent_slicing_min_size(): + vae = _SlicingDecodeVAE(slicing_latent_min_size=2) + z = torch.arange( + _LATENT_CHANNELS * 5 * 8 * 8, + dtype=torch.float32, + ).reshape(1, _LATENT_CHANNELS, 5, 8, 8) + + tiled_vae( + z, + vae, + tile_size=(64, 64), + tile_overlap=(0, 0), + temporal_size=12, + temporal_overlap=4, + encode=False, + ) + + assert vae.decode_min_sizes == [2] + assert vae.memory_states == [MemoryState.INITIALIZING, MemoryState.ACTIVE] + assert vae.slicing_latent_min_size == 2 + + wrapper = vae_mod.VideoAutoencoderKLWrapper.__new__( + vae_mod.VideoAutoencoderKLWrapper + ) + nn.Module.__init__(wrapper) + seedvr2_tiling = { + "enable_tiling": True, + "tile_size": (64, 64), + "tile_overlap": (0, 0), + "temporal_size": 8, + "temporal_overlap": 7, + } + + captured = {} + + def _fake_tiled_vae(latent, model, **kwargs): + captured.update(kwargs) + return torch.zeros(1, 3, 1, 16, 16) + + with patch.object(vae_mod, "tiled_vae", side_effect=_fake_tiled_vae): + wrapper.decode(torch.zeros(1, _LATENT_CHANNELS, 2, 2), seedvr2_tiling=seedvr2_tiling) + + assert captured["temporal_overlap"] == 7 + + +def _force_oom(*a, **k): + raise torch.cuda.OutOfMemoryError("forced OOM for dispatcher test") + + +def _make_vae(first_stage_model, latent_channels, latent_dim): + vae = sd_mod.VAE.__new__(sd_mod.VAE) + vae.first_stage_model = first_stage_model + vae.patcher = MagicMock() + vae.patcher.get_free_memory = MagicMock(return_value=8 * 1024 * 1024 * 1024) + vae.device = vae.output_device = torch.device("cpu") + vae.vae_dtype = torch.float32 + vae.disable_offload = True + vae.extra_1d_channel = None + vae.upscale_ratio = vae.downscale_ratio = 8 + vae.upscale_index_formula = vae.downscale_index_formula = None + vae.output_channels = 3 + vae.latent_channels = latent_channels + vae.latent_dim = latent_dim + vae.vae_output_dtype = lambda: torch.float32 + vae.spacial_compression_decode = lambda: 8 + vae.handles_tiling = isinstance(first_stage_model, seedvr_vae_mod.VideoAutoencoderKLWrapper) + vae.format_encoded = None + vae.process_input = lambda x: x + vae.process_output = lambda x: x + vae.throw_exception_if_invalid = lambda: None + vae.memory_used_decode = lambda *a, **k: 1 + return vae + + +def _dispatch(vae, samples, seedvr2_call, generic_call, patch_wrapper_decode): + mm = sd_mod.model_management + with ExitStack() as stack: + stack.enter_context(patch.object(mm, "raise_non_oom", lambda e: None)) + stack.enter_context(patch.object(mm, "load_models_gpu", lambda *a, **k: None)) + stack.enter_context(patch.object(mm, "soft_empty_cache", lambda: None)) + stack.enter_context(patch.object(sd_mod.VAE, "_decode_tiled_owned", seedvr2_call)) + stack.enter_context(patch.object(sd_mod.VAE, "decode_tiled_", generic_call)) + if patch_wrapper_decode: + stack.enter_context(patch.object( + seedvr_vae_mod.VideoAutoencoderKLWrapper, "decode", + side_effect=_force_oom)) + vae.decode(samples) + + +def test_4d_seedvr2_latent_routes_to_owned_decode_tiled(): + wrapper = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__( + seedvr_vae_mod.VideoAutoencoderKLWrapper) + vae = _make_vae(wrapper, latent_channels=_LATENT_CHANNELS, latent_dim=3) + seedvr2_call = MagicMock(return_value=torch.zeros(1, 3, 9, 64, 64)) + generic_call = MagicMock(return_value=torch.zeros(1, 3, 64, 64)) + _dispatch(vae, torch.zeros(1, _LATENT_CHANNELS * 3, 8, 8), seedvr2_call, generic_call, True) + assert seedvr2_call.call_count == 1 + assert generic_call.call_count == 0 + + +def test_4d_non_seedvr2_latent_still_routes_to_generic_decode_tiled(): + first_stage = MagicMock() + first_stage.decode = MagicMock(side_effect=_force_oom) + vae = _make_vae(first_stage, latent_channels=4, latent_dim=2) + seedvr2_call = MagicMock(return_value=torch.zeros(1, 3, 9, 64, 64)) + generic_call = MagicMock(return_value=torch.zeros(1, 3, 64, 64)) + _dispatch(vae, torch.zeros(1, 4, 8, 8), seedvr2_call, generic_call, False) + assert generic_call.call_count == 1 + assert seedvr2_call.call_count == 0 + + +def _populate_common_vae_attrs_fallback(vae): + vae.patcher = MagicMock() + vae.patcher.get_free_memory = MagicMock(return_value=8 * 1024 * 1024 * 1024) + vae.device = torch.device("cpu") + vae.output_device = torch.device("cpu") + vae.vae_dtype = torch.float32 + vae.disable_offload = True + vae.extra_1d_channel = None + vae.upscale_ratio = 8 + vae.upscale_index_formula = None + vae.output_channels = 3 + vae.latent_channels = _LATENT_CHANNELS + vae.latent_dim = 3 + vae.downscale_ratio = 8 + vae.downscale_index_formula = None + vae.not_video = False + vae.crop_input = False + vae.pad_channel_value = None + vae.handles_tiling = isinstance(vae.first_stage_model, seedvr_vae_mod.VideoAutoencoderKLWrapper) + vae.format_encoded = None + + vae.vae_output_dtype = lambda: torch.float32 + vae.spacial_compression_encode = lambda: 8 + vae.process_input = lambda x: x + vae.process_output = lambda x: x + vae.throw_exception_if_invalid = lambda: None + vae.memory_used_encode = lambda *a, **k: 1 + + +def _make_seedvr2_vae_fallback(): + vae = sd_mod.VAE.__new__(sd_mod.VAE) + wrapper = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__( + seedvr_vae_mod.VideoAutoencoderKLWrapper + ) + vae.first_stage_model = wrapper + _populate_common_vae_attrs_fallback(vae) + return vae + + +def _make_non_seedvr2_vae_fallback(): + vae = sd_mod.VAE.__new__(sd_mod.VAE) + vae.first_stage_model = MagicMock() + _populate_common_vae_attrs_fallback(vae) + return vae + + +def _force_regular_encode_oom(*args, **kwargs): + raise torch.cuda.OutOfMemoryError("forced OOM for dispatcher test") + + +def test_seedvr2_3d_routes_to_owned_encode_tiled_on_oom(): + vae = _make_seedvr2_vae_fallback() + pixel_samples = torch.zeros((1, 8, 64, 64, 3)) + + seedvr2_call = MagicMock(return_value=torch.zeros(1, _LATENT_CHANNELS, 2, 8, 8)) + generic_call = MagicMock(return_value=torch.zeros(1, _LATENT_CHANNELS, 2, 8, 8)) + + with patch.object(sd_mod.model_management, "raise_non_oom", + lambda e: None), \ + patch.object(sd_mod.model_management, "load_models_gpu", + lambda *a, **k: None), \ + patch.object(sd_mod.model_management, "soft_empty_cache", + lambda: None), \ + patch.object(seedvr_vae_mod.VideoAutoencoderKLWrapper, "encode", + side_effect=_force_regular_encode_oom), \ + patch.object(sd_mod.VAE, "_encode_tiled_owned", seedvr2_call), \ + patch.object(sd_mod.VAE, "encode_tiled_3d", generic_call): + vae.encode(pixel_samples) + + assert seedvr2_call.call_count == 1, ( + f"Expected _encode_tiled_owned to be called once for a SeedVR2 3D " + f"input under OOM fallback; got {seedvr2_call.call_count} calls." + ) + assert generic_call.call_count == 0, ( + f"encode_tiled_3d must NOT be called for a SeedVR2 input; got " + f"{generic_call.call_count} calls." + ) + + +def test_non_seedvr2_encode_tiled_3d_default_overlap_is_concrete(): + vae = _make_non_seedvr2_vae_fallback() + vae.downscale_ratio = (lambda a: max(1, a // 4), 8, 8) + vae.upscale_ratio = (lambda a: a * 4, 8, 8) + generic_call = MagicMock(return_value=torch.zeros(1, _LATENT_CHANNELS, 2, 8, 8)) + pixel_samples = torch.zeros((1, 8, 64, 64, 3)) + + with patch.object(sd_mod.model_management, "load_models_gpu", + lambda *a, **k: None), \ + patch.object(sd_mod.VAE, "encode_tiled_3d", generic_call): + vae.encode_tiled(pixel_samples) + + assert generic_call.call_args.kwargs["overlap"] == (1, 64, 64) From 89ecc5cf8c1992230ddab3ee67ef836b7c884442 Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Fri, 10 Jul 2026 11:58:22 +0300 Subject: [PATCH 10/37] [Partner Nodes] feat(Seedream): add widget to disable thinking (#14853) Signed-off-by: bigcat88 Co-authored-by: Daxiong (Lin) --- comfy_api_nodes/apis/bytedance.py | 5 +++++ comfy_api_nodes/nodes_bytedance.py | 21 +++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/comfy_api_nodes/apis/bytedance.py b/comfy_api_nodes/apis/bytedance.py index 76573304b..515e124ca 100644 --- a/comfy_api_nodes/apis/bytedance.py +++ b/comfy_api_nodes/apis/bytedance.py @@ -17,6 +17,10 @@ class Seedream4Options(BaseModel): max_images: int = Field(15) +class Seedream5OptimizePromptOptions(BaseModel): + thinking: Literal["auto", "enabled", "disabled"] = Field(...) + + class Seedream4TaskCreationRequest(BaseModel): model: str = Field(...) prompt: str = Field(...) @@ -28,6 +32,7 @@ class Seedream4TaskCreationRequest(BaseModel): sequential_image_generation_options: Seedream4Options | None = Field(Seedream4Options(max_images=15)) watermark: bool = Field(False) output_format: str | None = None + optimize_prompt_options: Seedream5OptimizePromptOptions | None = None class ImageTaskCreationResponse(BaseModel): diff --git a/comfy_api_nodes/nodes_bytedance.py b/comfy_api_nodes/nodes_bytedance.py index 043bc9526..a84399ad3 100644 --- a/comfy_api_nodes/nodes_bytedance.py +++ b/comfy_api_nodes/nodes_bytedance.py @@ -34,6 +34,7 @@ from comfy_api_nodes.apis.bytedance import ( SeedanceVirtualLibraryCreateAssetRequest, Seedream4Options, Seedream4TaskCreationRequest, + Seedream5OptimizePromptOptions, TaskAudioContent, TaskAudioContentUrl, TaskCreationResponse, @@ -875,6 +876,17 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode): tooltip='Whether to add an "AI generated" watermark to the image.', advanced=True, ), + IO.Boolean.Input( + "thinking", + default=True, + tooltip=( + "Enable the model's prompt-optimization reasoning ('thinking') for better adherence. " + "Can substantially increase generation time — notably on Seedream 5.0 Pro. " + "Can only be disabled for text-to-image (not when reference images are provided)." + ), + optional=True, + advanced=True, + ), ], outputs=[ IO.Image.Output(), @@ -920,6 +932,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode): model: dict, seed: int = 0, watermark: bool = False, + thinking: bool = True, ) -> IO.NodeOutput: validate_string(prompt, strip_whitespace=True, min_length=1) model_id = SEEDREAM_MODELS[model["model"]] @@ -979,6 +992,10 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode): raise ValueError( "The maximum number of generated images plus the number of reference images cannot exceed 15." ) + if not thinking and n_input_images > 0: + raise ValueError( + "'thinking' can only be disabled for text-to-image; enable it when using reference images." + ) reference_images_urls: list[str] = [] if image_tensors: @@ -992,6 +1009,9 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode): wait_label="Uploading reference images", ) + optimize_prompt_options = None + if n_input_images == 0: + optimize_prompt_options = Seedream5OptimizePromptOptions(thinking="enabled" if thinking else "disabled") response = await sync_op( cls, ApiEndpoint(path=BYTEPLUS_IMAGE_ENDPOINT, method="POST"), @@ -1005,6 +1025,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode): sequential_image_generation=None if is_pro else sequential_image_generation, sequential_image_generation_options=None if is_pro else Seedream4Options(max_images=max_images), watermark=watermark, + optimize_prompt_options=optimize_prompt_options, ), ) if len(response.data) == 1: From 206b9245dcc4e497bf18808adf7585b8ae7595ec Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Fri, 10 Jul 2026 12:33:32 +0300 Subject: [PATCH 11/37] [Partner Nodes] fix(Tencent): restore Tencent3DPartNode FBX output via staged generation (#14867) Signed-off-by: bigcat88 --- comfy_api_nodes/apis/hunyuan3d.py | 1 + comfy_api_nodes/nodes_hunyuan3d.py | 1 + 2 files changed, 2 insertions(+) diff --git a/comfy_api_nodes/apis/hunyuan3d.py b/comfy_api_nodes/apis/hunyuan3d.py index dad9bc2fa..91f630e81 100644 --- a/comfy_api_nodes/apis/hunyuan3d.py +++ b/comfy_api_nodes/apis/hunyuan3d.py @@ -77,6 +77,7 @@ class To3DUVTaskRequest(BaseModel): class To3DPartTaskRequest(BaseModel): File: TaskFile3DInput = Field(...) + EnableStagedGeneration: bool | None = Field(None) class TextureEditImageInfo(BaseModel): diff --git a/comfy_api_nodes/nodes_hunyuan3d.py b/comfy_api_nodes/nodes_hunyuan3d.py index fcd27b7fb..a9942476c 100644 --- a/comfy_api_nodes/nodes_hunyuan3d.py +++ b/comfy_api_nodes/nodes_hunyuan3d.py @@ -642,6 +642,7 @@ class Tencent3DPartNode(IO.ComfyNode): response_model=To3DProTaskCreateResponse, data=To3DPartTaskRequest( File=TaskFile3DInput(Type=file_format.upper(), Url=model_url), + EnableStagedGeneration=True, ), is_rate_limited=_is_tencent_rate_limited, ) From 1377a2f72925ed7a5518c1900ff71c6740217b0d Mon Sep 17 00:00:00 2001 From: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com> Date: Fri, 10 Jul 2026 18:31:20 +0800 Subject: [PATCH 12/37] Only auto-enable the ROCm comfy-kitchen Triton backend on matrix-core GPUs (#14869) #14862 auto-enables the comfy-kitchen Triton backend whenever torch.version.hip is set and Triton >= 3.7. The INT8 matmul kernels compile tl.dot to matrix-core instructions (WMMA on RDNA3+/gfx11xx-gfx12xx, MFMA on CDNA/gfx9xx); RDNA1/RDNA2 (gfx10xx) have neither, so the auto-enabled INT8 path hangs the GPU there (reported on RDNA2 + triton-windows 3.7.1: native and custom-node INT8 freeze until reset). Gate the automatic ROCm default on GPU architecture as well as Triton version so RDNA1/RDNA2 stay on the working eager fallback. Add --disable-triton-backend as an explicit override; --enable-triton-backend still force-enables on any arch. --- comfy/cli_args.py | 1 + comfy/quant_ops.py | 24 ++++++++++++++++++++++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 0d7df5e13..e2e0d97ec 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -92,6 +92,7 @@ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE" parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.") parser.add_argument("--supports-fp8-compute", action="store_true", help="ComfyUI will act like if the device supports fp8 compute.") parser.add_argument("--enable-triton-backend", action="store_true", help="ComfyUI will enable the use of Triton backend in comfy-kitchen. Is disabled at launch by default.") +parser.add_argument("--disable-triton-backend", action="store_true", help="Force-disable the comfy-kitchen Triton backend, overriding the automatic ROCm/AMD default and --enable-triton-backend.") class LatentPreviewMethod(enum.Enum): NoPreviews = "none" diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py index 91b3e4fe9..b1aabdc93 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -3,6 +3,22 @@ import logging from comfy.cli_args import args + +def _rocm_kitchen_arch_supported(): + """comfy-kitchen's INT8 Triton kernels compile tl.dot to matrix-core instructions. + RDNA3/3.5/4 (gfx11xx/gfx12xx) have WMMA and CDNA (gfx9xx) has MFMA; RDNA1/RDNA2 + (gfx10xx) have neither, so the INT8 path hangs the GPU there. Gates the automatic + ROCm default so those cards stay on the eager fallback (an explicit + --enable-triton-backend still forces it on any arch).""" + try: + arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0] + except Exception: + return False + if arch.startswith(("gfx11", "gfx12")): + return True + return arch in ("gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx950") + + try: import comfy_kitchen as ck from comfy_kitchen.tensor import ( @@ -26,9 +42,13 @@ try: logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.") # On ROCm/AMD the CUDA backend is unavailable, so Triton is the only accelerated - # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7: + # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7 AND a + # matrix-core GPU (RDNA3+ WMMA gfx11xx/gfx12xx, CDNA MFMA gfx9xx). RDNA1/RDNA2 + # (gfx10xx) have no WMMA -> the INT8 tl.dot path hangs the GPU, so they stay eager. # older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path. - if args.enable_triton_backend or torch.version.hip is not None: + if args.disable_triton_backend: + ck.registry.disable("triton") + elif args.enable_triton_backend or (torch.version.hip is not None and _rocm_kitchen_arch_supported()): try: import triton triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2]) From 94fa08223e611f1e95693fab90f8c00adf353ccc Mon Sep 17 00:00:00 2001 From: "Yousef R. Gamaleldin" <81116377+yousef-rafat@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:54:56 +0300 Subject: [PATCH 13/37] Save Text Node (CORE-176) (#14102) --- comfy_execution/jobs.py | 15 ++++++-- comfy_extras/nodes_text.py | 71 ++++++++++++++++++++++++++++++++++++++ nodes.py | 1 + 3 files changed, 84 insertions(+), 3 deletions(-) create mode 100644 comfy_extras/nodes_text.py diff --git a/comfy_execution/jobs.py b/comfy_execution/jobs.py index fa3ab0faf..f0ad59f86 100644 --- a/comfy_execution/jobs.py +++ b/comfy_execution/jobs.py @@ -56,6 +56,9 @@ PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio', '3d', 'text'}) # 3D file extensions for preview fallback (no dedicated media_type exists) THREE_D_EXTENSIONS = frozenset({'.obj', '.fbx', '.gltf', '.glb', '.usdz'}) +# Text file extensions for preview fallback (the formats SaveText can produce) +TEXT_EXTENSIONS = frozenset({'.txt', '.md', '.json'}) + def has_3d_extension(filename: str) -> bool: lower = filename.lower() @@ -143,9 +146,10 @@ def is_previewable(media_type: str, item: dict) -> bool: Maintains backwards compatibility with existing logic. Priority: - 1. media_type is 'images', 'video', 'audio', or '3d' + 1. media_type is 'images', 'video', 'audio', '3d', or 'text' 2. format field starts with 'video/' or 'audio/' 3. filename has a 3D extension (.obj, .fbx, .gltf, .glb, .usdz) + 4. filename has a text extension (.txt, .md, .json, ...) """ if media_type in PREVIEWABLE_MEDIA_TYPES: return True @@ -156,10 +160,12 @@ def is_previewable(media_type: str, item: dict) -> bool: if fmt and (fmt.startswith('video/') or fmt.startswith('audio/')): return True - # Check for 3D files by extension + # Check for 3D and text files by extension filename = item.get('filename', '').lower() if any(filename.endswith(ext) for ext in THREE_D_EXTENSIONS): return True + if any(filename.endswith(ext) for ext in TEXT_EXTENSIONS): + return True return False @@ -255,6 +261,10 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]: Preview priority (matching frontend): 1. type="output" with previewable media 2. Any previewable media + + Text content entries (strings under 'text') are preview-only metadata, + matching the frontend's METADATA_KEYS: they can serve as the fallback + preview but are not counted as outputs. """ count = 0 preview_output = None @@ -275,7 +285,6 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]: if normalized is None: # Not a 3D file string — check for text preview if media_type == 'text': - count += 1 if preview_output is None: if isinstance(item, tuple): text_value = item[0] if item else '' diff --git a/comfy_extras/nodes_text.py b/comfy_extras/nodes_text.py new file mode 100644 index 000000000..a485f5df8 --- /dev/null +++ b/comfy_extras/nodes_text.py @@ -0,0 +1,71 @@ +import os +import json +from typing_extensions import override +from comfy_api.latest import io, ComfyExtension, ui +import folder_paths + + +class SaveTextNode(io.ComfyNode): + """Save text content to .txt, .md, or .json.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="SaveText", + search_aliases=["save text", "write text", "export text"], + display_name="Save Text", + category="text", + description="Save text content to a file in the output directory.", + inputs=[ + io.String.Input("text", force_input=True), + io.String.Input("filename_prefix", default="ComfyUI"), + io.Combo.Input("format", options=["txt", "md", "json"], default="txt"), + ], + outputs=[io.String.Output(display_name="text")], + is_output_node=True, + ) + + @classmethod + def execute(cls, text, filename_prefix, format): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( + filename_prefix, + folder_paths.get_output_directory(), + 1, + 1, + ) + + file = f"{filename}_{counter:05}.{format}" + filepath = os.path.join(full_output_folder, file) + + if format == "json": + # tries to pretty print otherwise saves normally + try: + data = json.loads(text) + with open(filepath, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + except json.JSONDecodeError: + with open(filepath, "w", encoding="utf-8") as f: + f.write(text) + else: + with open(filepath, "w", encoding="utf-8") as f: + f.write(text) + + return io.NodeOutput( + text, + ui={ + "text": (text,), + "files": [ + ui.SavedResult(file, subfolder, io.FolderType.output) + ] + } + ) + +class TextExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + SaveTextNode + ] + +async def comfy_entrypoint() -> TextExtension: + return TextExtension() diff --git a/nodes.py b/nodes.py index 474e188fe..31602e582 100644 --- a/nodes.py +++ b/nodes.py @@ -2504,6 +2504,7 @@ async def init_builtin_extra_nodes(): "nodes_triposplat.py", "nodes_depth_anything_3.py", "nodes_seed.py", + "nodes_text.py", ] import_failed = [] From 8310b0e0dbc3361c69c70985ee63b73f1970449a Mon Sep 17 00:00:00 2001 From: Terry Jia Date: Fri, 10 Jul 2026 15:58:03 -0400 Subject: [PATCH 14/37] feat: add bboxes input to Create Bounding Boxes node (#14724) --- comfy_extras/nodes_bounding_boxes.py | 132 ++++++++++++++++++++++++++- 1 file changed, 129 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_bounding_boxes.py b/comfy_extras/nodes_bounding_boxes.py index 77cbf8649..de3709b91 100644 --- a/comfy_extras/nodes_bounding_boxes.py +++ b/comfy_extras/nodes_bounding_boxes.py @@ -1,3 +1,5 @@ +import json + import numpy as np import torch from PIL import Image, ImageDraw, ImageEnhance, ImageFont @@ -166,6 +168,111 @@ def boxes_to_regions(boxes, width: int, height: int) -> list: return regions +def normalize_incoming_boxes(bboxes) -> list: + if isinstance(bboxes, dict): + frame = [bboxes] + elif not isinstance(bboxes, list) or not bboxes: + frame = [] + elif isinstance(bboxes[0], dict): + frame = bboxes + else: + frame = bboxes[0] if isinstance(bboxes[0], list) else [] + boxes = [] + for box in frame: + if not isinstance(box, dict): + continue + norm = { + "x": box.get("x", 0), + "y": box.get("y", 0), + "width": box.get("width", 0), + "height": box.get("height", 0), + } + meta = box.get("metadata") + if isinstance(meta, dict): + norm["metadata"] = meta + boxes.append(norm) + return boxes + + +def _looks_like_element(box: dict) -> bool: + bbox = box.get("bbox") + return isinstance(bbox, (list, tuple)) and len(bbox) == 4 + + +def _looks_like_bbox(box: dict) -> bool: + return all(key in box for key in ("x", "y", "width", "height")) + + +def elements_to_boxes(elements: list, width: int, height: int) -> list: + boxes = [] + for element in elements: + if not isinstance(element, dict): + continue + bbox = element.get("bbox") + if not (isinstance(bbox, (list, tuple)) and len(bbox) == 4): + raise ValueError("bboxes element is missing a valid 'bbox' [ymin, xmin, ymax, xmax]") + try: + ymin, xmin, ymax, xmax = (float(v) / 1000.0 for v in bbox) + except (TypeError, ValueError): + raise ValueError("bboxes element 'bbox' must contain four numbers") + etype = "text" if element.get("type") == "text" else "obj" + boxes.append({ + "x": round(min(xmin, xmax) * width), + "y": round(min(ymin, ymax) * height), + "width": round(abs(xmax - xmin) * width), + "height": round(abs(ymax - ymin) * height), + "metadata": { + "type": etype, + "text": element.get("text", "") if etype == "text" else "", + "desc": element.get("desc", ""), + "palette": element.get("color_palette", []) or [], + }, + }) + return boxes + + +def boxes_from_input(data, width: int, height: int) -> list: + if data is None: + return [] + if isinstance(data, str): + text = data.strip() + if not text: + return [] + try: + data = json.loads(text) + except (ValueError, TypeError) as exc: + raise ValueError(f"bboxes string input is not valid JSON: {exc}") from exc + if isinstance(data, dict): + if _looks_like_element(data): + return elements_to_boxes([data], width, height) + if _looks_like_bbox(data): + return normalize_incoming_boxes(data) + raise ValueError( + "bboxes dict must be a bounding box (x, y, width, height) or an element (with a 'bbox')" + ) + if not isinstance(data, list): + raise ValueError( + "bboxes input must be bounding boxes, elements, or a JSON string, " + f"got {type(data).__name__}" + ) + if not data: + return [] + first = data[0] + if isinstance(first, list): + return normalize_incoming_boxes(data) + if isinstance(first, dict): + if _looks_like_element(first): + return elements_to_boxes(data, width, height) + if _looks_like_bbox(first): + return normalize_incoming_boxes(data) + raise ValueError( + "bboxes items must be bounding boxes (x, y, width, height) or elements (with a 'bbox')" + ) + raise ValueError( + f"bboxes list must contain bounding boxes or elements, got {type(first).__name__}" + ) + + def _norm_bbox(region: dict) -> list[int]: def grid(value: float) -> int: return max(0, min(1000, round(value * 1000))) @@ -217,29 +324,48 @@ class CreateBoundingBoxes(io.ComfyNode): optional=True, tooltip="Optional image used as background in the canvas and preview.", ), + io.MultiType.Input( + "bboxes", + [io.BoundingBox, io.Array, io.String], + optional=True, + tooltip="Bounding boxes, elements, or a JSON string to initialize the canvas. A new upstream value initializes the canvas; edits made on the canvas take priority and are kept until the upstream value changes again.", + ), io.Int.Input("width", default=1024, min=64, max=16384, step=16, tooltip="Width of the canvas and the pixel grid for the bounding boxes."), io.Int.Input("height", default=1024, min=64, max=16384, step=16, tooltip="Height of the canvas and the pixel grid for the bounding boxes."), editor_state, + io.BoundingBoxes.Input( + "last_incoming", + optional=True, + tooltip="Internal state managed by the canvas: the upstream bboxes value that last initialized it. Leave empty to re-initialize the canvas from the bboxes input on the next run.", + ), ], outputs=[ io.Image.Output(display_name="preview"), io.BoundingBox.Output(display_name="bboxes"), io.Array.Output(display_name="elements"), ], + is_output_node=True, is_experimental=True, ) @classmethod - def execute(cls, width, height, editor_state=None, background=None) -> io.NodeOutput: - regions = boxes_to_regions(editor_state, width, height) + def execute(cls, width, height, editor_state=None, last_incoming=None, background=None, bboxes=None) -> io.NodeOutput: + incoming = boxes_from_input(bboxes, width, height) + applied = last_incoming if isinstance(last_incoming, list) else [] + upstream_changed = bool(incoming) and incoming != applied + source = incoming if upstream_changed else (editor_state or []) + regions = boxes_to_regions(source, width, height) preview = render_preview(regions, width, height, _bg_from_image(background)) + ui = {"dims": [width, height]} + if incoming: + ui["input_bboxes"] = incoming return io.NodeOutput( preview, fractions_to_bbox_frame(regions, width, height), build_elements(regions), - ui={"dims": [width, height]}, + ui=ui, ) From 328144ce24c6ce4b979dd850027f95d4cfa8449a Mon Sep 17 00:00:00 2001 From: Terry Jia Date: Fri, 10 Jul 2026 16:03:34 -0400 Subject: [PATCH 15/37] CORE-329 feat: add Save 3D (Advanced) node family (#14701) --- comfy_extras/nodes_save_3d.py | 158 +++++++++++++++++++++++++++++++++- 1 file changed, 156 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_save_3d.py b/comfy_extras/nodes_save_3d.py index 1b6592bb2..7c524caa1 100644 --- a/comfy_extras/nodes_save_3d.py +++ b/comfy_extras/nodes_save_3d.py @@ -13,7 +13,7 @@ from typing_extensions import override import folder_paths from comfy.cli_args import args -from comfy_api.latest import ComfyExtension, IO, Types +from comfy_api.latest import ComfyExtension, IO, Types, UI def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False): @@ -406,10 +406,164 @@ class SaveGLB(IO.ComfyNode): return IO.NodeOutput(ui={"3d": results}) +def _save_file3d_to_output(model_3d: Types.File3D, filename_prefix: str) -> str: + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( + filename_prefix, folder_paths.get_output_directory() + ) + ext = model_3d.format or "glb" + saved_filename = f"{filename}_{counter:05}.{ext}" + model_3d.save_to(os.path.join(full_output_folder, saved_filename)) + return f"{subfolder}/{saved_filename}" if subfolder else saved_filename + + +def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput: + model_file = _save_file3d_to_output(model_3d, filename_prefix) + camera_info_input = kwargs.get("camera_info", None) + camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info'] + model_3d_info_input = kwargs.get("model_3d_info", None) + model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', []) + return IO.NodeOutput( + model_3d, + model_3d_info, + camera_info, + width, + height, + ui=UI.PreviewUI3DAdvanced(model_file, camera_info, model_3d_info), + ) + + +class Save3DAdvanced(IO.ComfyNode): + @classmethod + def define_schema(cls): + return IO.Schema( + node_id="Save3DAdvanced", + display_name="Save 3D (Advanced)", + search_aliases=["save 3d", "export 3d model", "save mesh advanced"], + category="3d", + is_experimental=True, + is_output_node=True, + inputs=[ + IO.MultiType.Input( + "model_3d", + types=[ + IO.File3DGLB, + IO.File3DGLTF, + IO.File3DFBX, + IO.File3DOBJ, + IO.File3DSTL, + IO.File3DUSDZ, + IO.File3DAny, + ], + tooltip="3D model file from an upstream 3D node.", + ), + IO.String.Input("filename_prefix", default="3d/ComfyUI"), + IO.Load3D.Input("viewport_state"), + IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True), + IO.Load3DCamera.Input("camera_info", optional=True, advanced=True), + IO.Int.Input("width", default=1024, min=1, max=4096, step=1), + IO.Int.Input("height", default=1024, min=1, max=4096, step=1), + ], + outputs=[ + IO.File3DAny.Output(display_name="model_3d"), + IO.Load3DModelInfo.Output(display_name="model_3d_info"), + IO.Load3DCamera.Output(display_name="camera_info"), + IO.Int.Output(display_name="width"), + IO.Int.Output(display_name="height"), + ], + ) + + @classmethod + def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput: + return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) + + +class SaveGaussianSplat(IO.ComfyNode): + @classmethod + def define_schema(cls): + return IO.Schema( + node_id="SaveGaussianSplat", + display_name="Save Splat", + search_aliases=["save splat", "save gaussian splat", "export gaussian", "export splat"], + category="3d", + is_experimental=True, + is_output_node=True, + inputs=[ + IO.MultiType.Input( + "model_3d", + types=[ + IO.File3DSplatAny, + IO.File3DPLY, + IO.File3DSPLAT, + IO.File3DSPZ, + IO.File3DKSPLAT, + ], + tooltip="A gaussian splat 3D file.", + ), + IO.String.Input("filename_prefix", default="3d/ComfyUI"), + IO.Load3D.Input("viewport_state"), + IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True), + IO.Load3DCamera.Input("camera_info", optional=True, advanced=True), + IO.Int.Input("width", default=1024, min=1, max=4096, step=1), + IO.Int.Input("height", default=1024, min=1, max=4096, step=1), + ], + outputs=[ + IO.File3DSplatAny.Output(display_name="model_3d"), + IO.Load3DModelInfo.Output(display_name="model_3d_info"), + IO.Load3DCamera.Output(display_name="camera_info"), + IO.Int.Output(display_name="width"), + IO.Int.Output(display_name="height"), + ], + ) + + @classmethod + def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput: + return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) + + +class SavePointCloud(IO.ComfyNode): + @classmethod + def define_schema(cls): + return IO.Schema( + node_id="SavePointCloud", + display_name="Save Point Cloud", + search_aliases=["save point cloud", "save pointcloud", "export point cloud"], + category="3d", + is_experimental=True, + is_output_node=True, + inputs=[ + IO.MultiType.Input( + "model_3d", + types=[ + IO.File3DPointCloudAny, + IO.File3DPLY, + ], + tooltip="Point cloud file (.ply)", + ), + IO.String.Input("filename_prefix", default="3d/ComfyUI"), + IO.Load3D.Input("viewport_state"), + IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True), + IO.Load3DCamera.Input("camera_info", optional=True, advanced=True), + IO.Int.Input("width", default=1024, min=1, max=4096, step=1), + IO.Int.Input("height", default=1024, min=1, max=4096, step=1), + ], + outputs=[ + IO.File3DPointCloudAny.Output(display_name="model_3d"), + IO.Load3DModelInfo.Output(display_name="model_3d_info"), + IO.Load3DCamera.Output(display_name="camera_info"), + IO.Int.Output(display_name="width"), + IO.Int.Output(display_name="height"), + ], + ) + + @classmethod + def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput: + return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) + + class Save3DExtension(ComfyExtension): @override async def get_node_list(self) -> list[type[IO.ComfyNode]]: - return [SaveGLB] + return [SaveGLB, Save3DAdvanced, SaveGaussianSplat, SavePointCloud] async def comfy_entrypoint() -> Save3DExtension: From 5976ee37cd0ff1a5d28d14228daf8e5710390836 Mon Sep 17 00:00:00 2001 From: Alexis Rolland Date: Sat, 11 Jul 2026 07:31:39 +0800 Subject: [PATCH 16/37] Bringing back the text node (#14870) --- comfy_extras/nodes_primitive.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/comfy_extras/nodes_primitive.py b/comfy_extras/nodes_primitive.py index 7f90daf14..35761863f 100644 --- a/comfy_extras/nodes_primitive.py +++ b/comfy_extras/nodes_primitive.py @@ -10,11 +10,10 @@ class String(io.ComfyNode): return io.Schema( node_id="PrimitiveString", search_aliases=["text", "string", "text box", "prompt"], - display_name="Text String (DEPRECATED)", + display_name="Text", category="utilities/primitive", inputs=[io.String.Input("value")], - outputs=[io.String.Output()], - is_deprecated=True + outputs=[io.String.Output()] ) @classmethod @@ -28,7 +27,7 @@ class StringMultiline(io.ComfyNode): return io.Schema( node_id="PrimitiveStringMultiline", search_aliases=["text", "string", "text multiline", "string multiline", "text box", "prompt"], - display_name="Input Text", + display_name="Text (Multiline)", category="utilities/primitive", essentials_category="Basics", inputs=[io.String.Input("value", multiline=True)], From 1f51e146a884462a28b2071c2268ce7c72550ce9 Mon Sep 17 00:00:00 2001 From: Alexis Rolland Date: Sat, 11 Jul 2026 07:32:53 +0800 Subject: [PATCH 17/37] chore: Update preview nodes (#14871) --- comfy_extras/nodes_audio.py | 1 + comfy_extras/nodes_load_3d.py | 4 ++++ comfy_extras/nodes_mask.py | 5 +++-- comfy_extras/nodes_preview_any.py | 1 + nodes.py | 1 + 5 files changed, 10 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py index 6adcc95fa..4ac5ced53 100644 --- a/comfy_extras/nodes_audio.py +++ b/comfy_extras/nodes_audio.py @@ -298,6 +298,7 @@ class PreviewAudio(IO.ComfyNode): search_aliases=["play audio"], display_name="Preview Audio", category="audio", + description="Preview the audio without saving it to the ComfyUI output directory.", inputs=[ IO.Audio.Input("audio"), ], diff --git a/comfy_extras/nodes_load_3d.py b/comfy_extras/nodes_load_3d.py index 6ef9a1ca3..a9df557c2 100644 --- a/comfy_extras/nodes_load_3d.py +++ b/comfy_extras/nodes_load_3d.py @@ -92,6 +92,7 @@ class Preview3D(IO.ComfyNode): search_aliases=["view mesh", "3d viewer"], display_name="Preview 3D & Animation", category="3d", + description="Preview a 3D model file without saving it to the ComfyUI output directory.", is_experimental=True, is_output_node=True, inputs=[ @@ -136,6 +137,7 @@ class Preview3DAdvanced(IO.ComfyNode): display_name="Preview 3D (Advanced)", search_aliases=["preview 3d", "3d viewer", "view mesh", "frame 3d", "3d camera output"], category="3d", + description="Preview a 3D model file without saving it to the ComfyUI output directory.", is_experimental=True, is_output_node=True, inputs=[ @@ -193,6 +195,7 @@ class PreviewGaussianSplat(IO.ComfyNode): node_id="PreviewGaussianSplat", display_name="Preview Splat", category="3d", + description="Preview a gaussian splat 3D file without saving it to the ComfyUI output directory.", is_experimental=True, is_output_node=True, search_aliases=[ @@ -261,6 +264,7 @@ class PreviewPointCloud(IO.ComfyNode): node_id="PreviewPointCloud", display_name="Preview Point Cloud", category="3d", + description="Preview a point cloud 3D file without saving it to the ComfyUI output directory.", is_experimental=True, is_output_node=True, search_aliases=[ diff --git a/comfy_extras/nodes_mask.py b/comfy_extras/nodes_mask.py index 76af338de..3fae7221f 100644 --- a/comfy_extras/nodes_mask.py +++ b/comfy_extras/nodes_mask.py @@ -419,17 +419,18 @@ class MaskPreview(IO.ComfyNode): search_aliases=["show mask", "view mask", "inspect mask", "debug mask"], display_name="Preview Mask", category="image/mask", - description="Saves the input images to your ComfyUI output directory.", + description="Preview the masks without saving them to the ComfyUI output directory.", inputs=[ IO.Mask.Input("mask"), ], hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo], is_output_node=True, + outputs=[IO.Mask.Output(display_name="mask")] ) @classmethod def execute(cls, mask, filename_prefix="ComfyUI") -> IO.NodeOutput: - return IO.NodeOutput(ui=UI.PreviewMask(mask)) + return IO.NodeOutput(mask, ui=UI.PreviewMask(mask)) class MaskExtension(ComfyExtension): diff --git a/comfy_extras/nodes_preview_any.py b/comfy_extras/nodes_preview_any.py index 1070a69d0..d985f3287 100644 --- a/comfy_extras/nodes_preview_any.py +++ b/comfy_extras/nodes_preview_any.py @@ -18,6 +18,7 @@ class PreviewAny(): CATEGORY = "utilities" SEARCH_ALIASES = ["show output", "inspect", "debug", "print value", "show text"] + DESCRIPTION = "Preview any input value as text." def main(self, source=None): torch.set_printoptions(edgeitems=6) diff --git a/nodes.py b/nodes.py index 31602e582..883258bd1 100644 --- a/nodes.py +++ b/nodes.py @@ -1709,6 +1709,7 @@ class PreviewImage(SaveImage): self.compress_level = 1 SEARCH_ALIASES = ["preview", "preview image", "show image", "view image", "display image", "image viewer"] + DESCRIPTION = "Preview the images without saving them to the ComfyUI output directory." @classmethod def INPUT_TYPES(s): From 92ddf07ba14711cb579ab090846e0d51289c0619 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Fri, 10 Jul 2026 16:54:28 -0700 Subject: [PATCH 18/37] Try to fix some issues with the seedvr VAE. (#14877) --- comfy/ldm/seedvr/vae.py | 20 ++++++------- tests-unit/comfy_test/test_seedvr2_dtype.py | 29 +++++++++++++++++++ .../comfy_test/test_seedvr2_vae_tiled.py | 25 ++++++++++++++++ 3 files changed, 63 insertions(+), 11 deletions(-) diff --git a/comfy/ldm/seedvr/vae.py b/comfy/ldm/seedvr/vae.py index c9f430184..7a8070b65 100644 --- a/comfy/ldm/seedvr/vae.py +++ b/comfy/ldm/seedvr/vae.py @@ -30,7 +30,7 @@ from enum import Enum import logging import comfy.model_management import comfy.ops -ops = comfy.ops.disable_weight_init +ops = comfy.ops.manual_cast def _seedvr2_temporal_slicing_min_size(temporal_size, temporal_overlap, temporal_scale=1): @@ -103,11 +103,10 @@ def tiled_vae( storage_device = vae_model.device result = None count = None - def run_temporal_chunks(spatial_tile, model=vae_model, device=storage_device): - device = torch.device(device) - t_chunk = spatial_tile.to(device=device, dtype=next(model.parameters()).dtype, non_blocking=True).contiguous() + def run_temporal_chunks(spatial_tile, model=vae_model): + t_chunk = spatial_tile.contiguous() old_device = getattr(model, "device", None) - model.device = device + model.device = t_chunk.device old_slicing_min_size = getattr(model, slicing_attr, None) if old_slicing_min_size is not None and slicing_min_size is not None: if slicing_min_size <= 0: @@ -397,7 +396,7 @@ class Attention(nn.Module): def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor: input_dtype = x.dtype - if isinstance(norm_layer, (ops.LayerNorm, ops.RMSNorm)): + if isinstance(norm_layer, (nn.LayerNorm, nn.RMSNorm)): if x.ndim == 4: x = x.permute(0, 2, 3, 1) x = norm_layer(x) @@ -408,14 +407,14 @@ def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor: x = norm_layer(x) x = x.permute(0, 4, 1, 2, 3) return x.to(input_dtype) - if isinstance(norm_layer, (ops.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)): + if isinstance(norm_layer, (nn.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)): if x.ndim <= 4: return norm_layer(x).to(input_dtype) if x.ndim == 5: b, c, t, h, w = x.shape x = x.transpose(1, 2).reshape(b * t, c, h, w) memory_occupy = x.numel() * x.element_size() / 1024**3 - if isinstance(norm_layer, ops.GroupNorm) and memory_occupy > get_norm_limit(): + if isinstance(norm_layer, nn.GroupNorm) and memory_occupy > get_norm_limit(): num_chunks = min(BYTEDANCE_GN_CHUNKS_FP16 if x.element_size() == 2 else BYTEDANCE_GN_CHUNKS_FP32, norm_layer.num_groups) if norm_layer.num_groups % num_chunks != 0: raise ValueError( @@ -423,9 +422,9 @@ def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor: ) num_groups_per_chunk = norm_layer.num_groups // num_chunks + weights = comfy.ops.cast_to_input(norm_layer.weight, x).chunk(num_chunks, dim=0) + biases = comfy.ops.cast_to_input(norm_layer.bias, x).chunk(num_chunks, dim=0) x = list(x.chunk(num_chunks, dim=1)) - weights = norm_layer.weight.chunk(num_chunks, dim=0) - biases = norm_layer.bias.chunk(num_chunks, dim=0) for i, (w, bias) in enumerate(zip(weights, biases)): x[i] = F.group_norm(x[i], num_groups_per_chunk, w, bias, norm_layer.eps) x[i] = x[i].to(input_dtype) @@ -1459,7 +1458,6 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL): def _encode_with_raw_latent(self, x): if x.ndim == 4: x = x.unsqueeze(2) - x = x.to(dtype=next(self.parameters()).dtype) self.device = x.device p = super().encode(x) z = p.squeeze(2) diff --git a/tests-unit/comfy_test/test_seedvr2_dtype.py b/tests-unit/comfy_test/test_seedvr2_dtype.py index 8e08b6dde..d743cc848 100644 --- a/tests-unit/comfy_test/test_seedvr2_dtype.py +++ b/tests-unit/comfy_test/test_seedvr2_dtype.py @@ -1,4 +1,5 @@ import torch +import torch.nn as nn from comfy.cli_args import args as cli_args @@ -48,3 +49,31 @@ def test_seedvr2_vae_decode_memory_covers_full_frame_lab_transfer(): assert estimate == 101 * 960 * 1280 * 160 assert estimate > 15 * 1024 ** 3 assert estimate > old_estimate * 100 + + +def test_seedvr2_vae_encode_preserves_compute_dtype(monkeypatch): + wrapper = seedvr_vae.VideoAutoencoderKLWrapper.__new__(seedvr_vae.VideoAutoencoderKLWrapper) + nn.Module.__init__(wrapper) + wrapper._dummy = nn.Parameter(torch.empty(1, dtype=torch.float16)) + input_dtype = None + + def encode(self, x): + nonlocal input_dtype + input_dtype = x.dtype + return x + + monkeypatch.setattr(seedvr_vae.VideoAutoencoderKL, "encode", encode) + + x = torch.zeros((1, 3, 1, 8, 8), dtype=torch.float32) + wrapper._encode_with_raw_latent(x) + + assert input_dtype == torch.float32 + + +def test_seedvr2_vae_ops_cast_weights_to_compute_dtype(): + attention = seedvr_vae.Attention(query_dim=4, heads=1, dim_head=4).to(torch.float16) + hidden_states = torch.zeros((1, 2, 4), dtype=torch.float32) + + output = attention(hidden_states) + + assert output.dtype == torch.float32 diff --git a/tests-unit/comfy_test/test_seedvr2_vae_tiled.py b/tests-unit/comfy_test/test_seedvr2_vae_tiled.py index a2866b609..d64f51918 100644 --- a/tests-unit/comfy_test/test_seedvr2_vae_tiled.py +++ b/tests-unit/comfy_test/test_seedvr2_vae_tiled.py @@ -122,6 +122,31 @@ def test_tiled_vae_encode_uses_tensor_return_without_indexing(): assert tuple(out.shape) == (2, _LATENT_CHANNELS, 1, 8, 8) +def test_tiled_vae_preserves_compute_dtype_with_different_parameter_dtype(): + class DummyVAE(nn.Module): + spatial_downsample_factor = 8 + temporal_downsample_factor = 4 + slicing_sample_min_size = 8 + + def __init__(self): + super().__init__() + self.device = torch.device("cpu") + self._dummy = nn.Parameter(torch.zeros(1, dtype=torch.float16)) + self.input_dtype = None + + def encode(self, t_chunk): + self.input_dtype = t_chunk.dtype + b, _, _, h, w = t_chunk.shape + return torch.ones((b, _LATENT_CHANNELS, 1, h // 8, w // 8), dtype=t_chunk.dtype) + + vae = DummyVAE() + x = torch.zeros((1, 3, 1, 64, 64), dtype=torch.float32) + + tiled_vae(x, vae, tile_size=(64, 64), tile_overlap=(16, 16), encode=True) + + assert vae.input_dtype == torch.float32 + + def test_tiled_vae_preserves_input_dtype_on_single_tile(): class FloatOutputVAEModel(torch.nn.Module): def __init__(self): From f3a36e74844893f32f77f22d249d08862805d8f4 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Fri, 10 Jul 2026 18:37:59 -0700 Subject: [PATCH 19/37] Temporarily disable auto enabling triton by default on AMD. (#14878) I get freezing issues on my test machine. --- comfy/quant_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py index b1aabdc93..15f9b1fdb 100644 --- a/comfy/quant_ops.py +++ b/comfy/quant_ops.py @@ -48,7 +48,7 @@ try: # older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path. if args.disable_triton_backend: ck.registry.disable("triton") - elif args.enable_triton_backend or (torch.version.hip is not None and _rocm_kitchen_arch_supported()): + elif args.enable_triton_backend: # or (torch.version.hip is not None and _rocm_kitchen_arch_supported()): try: import triton triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2]) From 69ea58697bb2f05124f5dc7e00ad111f7cfff645 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Sat, 11 Jul 2026 17:16:40 -0700 Subject: [PATCH 20/37] Try to fix flash attention related issue on AMD. (#14880) --- comfy/ldm/modules/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 2411aff5c..e6500cff4 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -709,7 +709,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape return out try: - @torch.library.custom_op("flash_attention::flash_attn", mutates_args=()) + @torch.library.custom_op("comfy::flash_attn", mutates_args=()) def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor: softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale From 8b099de36acd81acd1afa3b5442951dc847e0a52 Mon Sep 17 00:00:00 2001 From: Gustavo Schneiter Date: Sun, 12 Jul 2026 01:58:25 -0300 Subject: [PATCH 21/37] Fix SaveVideo description: says images, saves video (#14885) --- comfy_extras/nodes_video.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py index d3acc9ad0..3bfd00be4 100644 --- a/comfy_extras/nodes_video.py +++ b/comfy_extras/nodes_video.py @@ -81,7 +81,7 @@ class SaveVideo(io.ComfyNode): display_name="Save Video", category="video", essentials_category="Basics", - description="Saves the input images to your ComfyUI output directory.", + description="Saves the input videos to your ComfyUI output directory.", inputs=[ io.Video.Input("video", tooltip="The video to save."), io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."), From 917faef771a2fd2f14f44af94f17da3d0b2803a3 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Sun, 12 Jul 2026 09:43:30 -0700 Subject: [PATCH 22/37] Support PID 1.5 models. (#14894) --- comfy/ldm/pixeldit/model.py | 4 ++ comfy/ldm/pixeldit/pid.py | 64 +++++++++++++++---- comfy/model_detection.py | 37 ++++++++++- tests-unit/comfy_test/model_detection_test.py | 52 +++++++++++++++ 4 files changed, 140 insertions(+), 17 deletions(-) diff --git a/comfy/ldm/pixeldit/model.py b/comfy/ldm/pixeldit/model.py index b044b9b29..3b30b9226 100644 --- a/comfy/ldm/pixeldit/model.py +++ b/comfy/ldm/pixeldit/model.py @@ -197,6 +197,9 @@ class PixDiT_T2I(nn.Module): """Hook for subclasses to inject per-block state into the patch stream (e.g. PiD's LQ gate).""" return s + def _pre_pixel_blocks(self, s, **kwargs): + return s + def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs): H_orig, W_orig = x.shape[2], x.shape[3] x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size)) @@ -226,6 +229,7 @@ class PixDiT_T2I(nn.Module): s, y_emb = blk(s, y_emb, condition, pos_img, pos_txt, None, transformer_options=transformer_options) s = F.silu(t_emb + s) + s = self._pre_pixel_blocks(s, **kwargs) s_cond = s.view(B * L, self.hidden_size) x_pixels = self.pixel_embedder(x, patch_size=self.patch_size) for blk in self.pixel_blocks: diff --git a/comfy/ldm/pixeldit/pid.py b/comfy/ldm/pixeldit/pid.py index 21b73907a..8590408d9 100644 --- a/comfy/ldm/pixeldit/pid.py +++ b/comfy/ldm/pixeldit/pid.py @@ -13,15 +13,15 @@ from .model import PixDiT_T2I from .modules import precompute_freqs_cis_2d -class SigmaAwareGatePerTokenPerDim(nn.Module): +class SigmaAwareGate(nn.Module): """gate = sigmoid(content_proj(cat[x, lq]) - exp(log_alpha) * sigma); out = x + gate * lq. Trained init gives ~0.88 gate at sigma=0, ~0.05 at sigma=1. """ - def __init__(self, dim: int, dtype=None, device=None, operations=None): + def __init__(self, dim: int, per_token: bool = False, dtype=None, device=None, operations=None): super().__init__() - self.content_proj = operations.Linear(dim * 2, dim, dtype=dtype, device=device) + self.content_proj = operations.Linear(dim * 2, 1 if per_token else dim, dtype=dtype, device=device) self.log_alpha = nn.Parameter(torch.empty((), dtype=dtype, device=device)) def forward(self, x: torch.Tensor, lq: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: @@ -36,15 +36,15 @@ class SigmaAwareGatePerTokenPerDim(nn.Module): class ResBlock(nn.Module): """Pre-activation ResNet block: GN -> SiLU -> Conv -> GN -> SiLU -> Conv + skip.""" - def __init__(self, channels: int, num_groups: int = 4, dtype=None, device=None, operations=None): + def __init__(self, channels: int, num_groups: int = 4, conv_padding_mode: str = "zeros", dtype=None, device=None, operations=None): super().__init__() self.block = nn.Sequential( operations.GroupNorm(num_groups, channels, dtype=dtype, device=device), nn.SiLU(), - operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device), + operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device), operations.GroupNorm(num_groups, channels, dtype=dtype, device=device), nn.SiLU(), - operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device), + operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device), ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -62,9 +62,13 @@ class LQProjection2D(nn.Module): patch_size: int = 16, sr_scale: int = 4, latent_spatial_down_factor: int = 8, + latent_unpatchify_factor: int = 1, num_res_blocks: int = 4, num_outputs: int = 7, interval: int = 2, + conv_padding_mode: str = "zeros", + gate_per_token: bool = False, + pit_output: bool = False, dtype=None, device=None, operations=None, ): super().__init__() @@ -74,34 +78,38 @@ class LQProjection2D(nn.Module): self.patch_size = patch_size self.sr_scale = sr_scale self.latent_spatial_down_factor = latent_spatial_down_factor + self.latent_unpatchify_factor = latent_unpatchify_factor self.num_outputs = num_outputs self.interval = interval - z_to_patch_ratio = (sr_scale * latent_spatial_down_factor) / patch_size + effective_latent_channels = latent_channels // (latent_unpatchify_factor * latent_unpatchify_factor) + effective_spatial_down_factor = latent_spatial_down_factor // latent_unpatchify_factor + z_to_patch_ratio = (sr_scale * effective_spatial_down_factor) / patch_size self.z_to_patch_ratio = z_to_patch_ratio if z_to_patch_ratio >= 1: self.latent_fold_factor = 0 - latent_proj_in_ch = latent_channels + latent_proj_in_ch = effective_latent_channels else: fold_factor = int(1 / z_to_patch_ratio) assert fold_factor * z_to_patch_ratio == 1.0 self.latent_fold_factor = fold_factor - latent_proj_in_ch = latent_channels * fold_factor * fold_factor + latent_proj_in_ch = effective_latent_channels * fold_factor * fold_factor layers = [ - operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device), + operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device), nn.SiLU(), - operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device), + operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device), ] for _ in range(num_res_blocks): - layers.append(ResBlock(hidden_dim, dtype=dtype, device=device, operations=operations)) + layers.append(ResBlock(hidden_dim, conv_padding_mode=conv_padding_mode, dtype=dtype, device=device, operations=operations)) self.latent_proj = nn.Sequential(*layers) self.output_heads = nn.ModuleList( [operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) for _ in range(num_outputs)] ) + self.pit_head = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) if pit_output else None self.gate_modules = nn.ModuleList( - [SigmaAwareGatePerTokenPerDim(out_dim, dtype=dtype, device=device, operations=operations) + [SigmaAwareGate(out_dim, per_token=gate_per_token, dtype=dtype, device=device, operations=operations) for _ in range(num_outputs)] ) @@ -115,6 +123,11 @@ class LQProjection2D(nn.Module): return self.gate_modules[out_idx](x, lq_feature, sigma) def _align_latent_to_patch_grid(self, lq_latent: torch.Tensor, pH: int, pW: int) -> torch.Tensor: + f = self.latent_unpatchify_factor + if f > 1: + B, C, H, W = lq_latent.shape + lq_latent = lq_latent.reshape(B, C // (f * f), f, f, H, W) + lq_latent = lq_latent.permute(0, 1, 4, 2, 5, 3).reshape(B, C // (f * f), H * f, W * f) B, z_dim = lq_latent.shape[:2] if self.z_to_patch_ratio >= 1: if lq_latent.shape[2] != pH or lq_latent.shape[3] != pW: @@ -134,7 +147,10 @@ class LQProjection2D(nn.Module): feat = self._align_latent_to_patch_grid(lq_latent, target_pH, target_pW) B, C, H, W = feat.shape tokens = feat.permute(0, 2, 3, 1).contiguous().view(B, H * W, C) - return [head(tokens) for head in self.output_heads] + outputs = [head(tokens) for head in self.output_heads] + if self.pit_head is not None: + outputs.append(self.pit_head(tokens)) + return outputs class PidNet(PixDiT_T2I): @@ -148,6 +164,10 @@ class PidNet(PixDiT_T2I): lq_interval: int = 2, sr_scale: int = 4, latent_spatial_down_factor: int = 8, + lq_latent_unpatchify_factor: int = 1, + lq_conv_padding_mode: str = "zeros", + lq_gate_per_token: bool = False, + pit_lq_inject: bool = False, rope_ref_h: int = 1024, # NTK ref resolution in PIXEL units: 1024px / patch=16 -> grid_ref=64. rope_ref_w: int = 1024, image_model=None, @@ -165,6 +185,8 @@ class PidNet(PixDiT_T2I): for blk in self.pixel_blocks: blk._rope_fn = _pit_rope_fn + self.pit_lq_inject = pit_lq_inject + num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval self.lq_proj = LQProjection2D( latent_channels=lq_latent_channels, @@ -173,13 +195,20 @@ class PidNet(PixDiT_T2I): patch_size=self.patch_size, sr_scale=sr_scale, latent_spatial_down_factor=latent_spatial_down_factor, + latent_unpatchify_factor=lq_latent_unpatchify_factor, num_res_blocks=lq_num_res_blocks, num_outputs=num_lq_outputs, interval=lq_interval, + conv_padding_mode=lq_conv_padding_mode, + gate_per_token=lq_gate_per_token, + pit_output=pit_lq_inject, dtype=dtype, device=device, operations=operations, ) + self.pit_lq_gate = SigmaAwareGate( + self.hidden_size, per_token=lq_gate_per_token, dtype=dtype, device=device, operations=operations + ) if pit_lq_inject else None def _fetch_patch_pos(self, height, width, device, dtype, **rope_opts): return precompute_freqs_cis_2d( @@ -197,6 +226,11 @@ class PidNet(PixDiT_T2I): return s return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx) + def _pre_pixel_blocks(self, s, pid_pit_lq_feature=None, pid_degrade_sigma=None, **kwargs): + if pid_pit_lq_feature is None: + return s + return self.pit_lq_gate(s, pid_pit_lq_feature, pid_degrade_sigma) + def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs): if lq_latent is None: raise ValueError("PidNet requires lq_latent — attach via PiDConditioning") @@ -216,12 +250,14 @@ class PidNet(PixDiT_T2I): degrade_sigma = degrade_sigma.expand(B).contiguous() lq_features = self.lq_proj(lq_latent=lq_latent.to(x), target_pH=Hs, target_pW=Ws) + pit_lq_feature = lq_features.pop() if self.pit_lq_inject else None return super()._forward( x, timesteps, context=context, attention_mask=attention_mask, transformer_options=transformer_options, pid_lq_features=lq_features, + pid_pit_lq_feature=pit_lq_feature, pid_degrade_sigma=degrade_sigma, **kwargs, ) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 174bc77cc..70c8625e3 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -470,15 +470,46 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): # PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I. _lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix) if _lq_w_key in state_dict_keys: - in_ch = int(state_dict[_lq_w_key].shape[1]) + latent_proj_in_channels = int(state_dict[_lq_w_key].shape[1]) + hidden_dim = int(state_dict[_lq_w_key].shape[0]) _gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix) num_gates = len({k[len(_gate_prefix):].split('.')[0] for k in state_dict_keys if k.startswith(_gate_prefix)}) + pid_v1_5 = '{}lq_proj.pit_head.weight'.format(key_prefix) in state_dict_keys dit_config = {"image_model": "pid", - "lq_latent_channels": in_ch, - "latent_spatial_down_factor": 16 if in_ch >= 64 else 8} + "lq_hidden_dim": hidden_dim} if num_gates > 0: dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates + if pid_v1_5: + pid_v1_5_variants = { + 16: { # Flux and QwenImage + "lq_latent_channels": 16, + "latent_spatial_down_factor": 8, + "lq_latent_unpatchify_factor": 1, + }, + 32: { # Flux2 after 2x latent unpatchify + "lq_latent_channels": 128, + "latent_spatial_down_factor": 16, + "lq_latent_unpatchify_factor": 2, + }, + } + variant = pid_v1_5_variants.get(latent_proj_in_channels) + if variant is None: + raise ValueError(f"Unsupported PiD v1.5 latent projection with {latent_proj_in_channels} input channels") + gate_weight = state_dict['{}lq_proj.gate_modules.0.content_proj.weight'.format(key_prefix)] + dit_config.update(variant) + dit_config.update({ + "lq_conv_padding_mode": "replicate", + "lq_gate_per_token": gate_weight.shape[0] == 1, + "pit_lq_inject": True, + "rope_ref_h": 2048, + "rope_ref_w": 2048, + }) + else: + dit_config.update({ + "lq_latent_channels": latent_proj_in_channels, + "latent_spatial_down_factor": 16 if latent_proj_in_channels >= 64 else 8, + }) return dit_config if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys: # PixelDiT T2I diff --git a/tests-unit/comfy_test/model_detection_test.py b/tests-unit/comfy_test/model_detection_test.py index 6e7d71f79..7c5b271c5 100644 --- a/tests-unit/comfy_test/model_detection_test.py +++ b/tests-unit/comfy_test/model_detection_test.py @@ -97,6 +97,21 @@ def _make_seedvr2_3b_shared_mm_sd(): } +def _make_pid_v1_5_sd(latent_proj_channels=16): + sd = { + "pixel_embedder.proj.weight": torch.empty(16, 3, device="meta"), + "lq_proj.latent_proj.0.weight": torch.empty(1024, latent_proj_channels, 3, 3, device="meta"), + "lq_proj.pit_head.weight": torch.empty(1536, 1024, device="meta"), + "lq_proj.gate_modules.0.content_proj.weight": torch.empty(1, 3072, device="meta"), + "pixel_blocks.0.attn.q_norm.weight": torch.empty(72, device="meta"), + "pixel_blocks.0.adaLN_modulation.0.weight": torch.empty(24576, 1536, device="meta"), + "pixel_blocks.0.adaLN_modulation.0.bias": torch.empty(24576, device="meta"), + } + for i in range(7): + sd[f"lq_proj.gate_modules.{i}.log_alpha"] = torch.empty((), device="meta") + return sd + + def _add_model_diffusion_prefix(sd): return {f"model.diffusion_model.{k}": v for k, v in sd.items()} @@ -206,6 +221,43 @@ class TestModelDetection: assert type(model_config_from_unet(sd, "model.diffusion_model.")).__name__ == "SeedVR2" + def test_pid_v1_5_detection(self): + sd = _make_pid_v1_5_sd() + unet_config = detect_unet_config(sd, "") + + assert unet_config == { + "image_model": "pid", + "lq_latent_channels": 16, + "lq_hidden_dim": 1024, + "latent_spatial_down_factor": 8, + "lq_interval": 2, + "lq_latent_unpatchify_factor": 1, + "lq_conv_padding_mode": "replicate", + "lq_gate_per_token": True, + "pit_lq_inject": True, + "rope_ref_h": 2048, + "rope_ref_w": 2048, + } + assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "PiD" + + def test_pid_v1_5_flux2_detection(self): + unet_config = detect_unet_config(_make_pid_v1_5_sd(latent_proj_channels=32), "") + + assert unet_config["lq_latent_channels"] == 128 + assert unet_config["latent_spatial_down_factor"] == 16 + assert unet_config["lq_latent_unpatchify_factor"] == 2 + + def test_pid_v1_5_pixel_adaln_conversion(self): + sd = _make_pid_v1_5_sd() + model_config = model_config_from_unet_config(detect_unet_config(sd, ""), sd) + processed = model_config.process_unet_state_dict(sd) + + assert processed["pixel_blocks.0.attn.q_norm.weight"].shape == (72,) + assert processed["pixel_blocks.0.adaLN_modulation_msa.weight"].shape == (12288, 1536) + assert processed["pixel_blocks.0.adaLN_modulation_mlp.weight"].shape == (12288, 1536) + assert processed["pixel_blocks.0.adaLN_modulation_msa.bias"].shape == (12288,) + assert processed["pixel_blocks.0.adaLN_modulation_mlp.bias"].shape == (12288,) + def test_unet_config_and_required_keys_combination_is_unique(self): """Each model in the registry must have a unique combination of ``unet_config`` and ``required_keys``. If two models share the same From b58f829b570d52fbbd41ebb00c6fe02b8755ec75 Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:38:35 +0300 Subject: [PATCH 23/37] [Partner Nodes] feat(client): send ComfyUI Core version in request headers (#14910) Signed-off-by: bigcat88 --- comfy_api_nodes/util/_helpers.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/comfy_api_nodes/util/_helpers.py b/comfy_api_nodes/util/_helpers.py index 7eb1ec664..acab10d95 100644 --- a/comfy_api_nodes/util/_helpers.py +++ b/comfy_api_nodes/util/_helpers.py @@ -15,6 +15,7 @@ from comfy.comfy_api_env import normalize_comfy_api_base from comfy.deploy_environment import get_deploy_environment from comfy.model_management import processing_interrupted from comfy_api.latest import IO +from comfyui_version import __version__ as comfyui_version from .common_exceptions import ProcessingInterrupted @@ -60,6 +61,7 @@ def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]: **get_auth_header(node_cls), "Comfy-Env": get_deploy_environment(), "Comfy-Usage-Source": get_usage_source(node_cls), + "Comfy-Core-Version": comfyui_version, } From 8deaa4d911497f93bbd434a3821efab396f6981f Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:53:37 +0300 Subject: [PATCH 24/37] fix(image): correct HLG inverse-OETF clamp in hlg_to_linear (#14762) Signed-off-by: bigcat88 Co-authored-by: Alexis Rolland --- comfy_extras/nodes_images.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index fe1937ba5..4d7b37200 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -891,10 +891,11 @@ def hlg_to_linear(t: torch.Tensor) -> torch.Tensor: return torch.cat([hlg_to_linear(rgb), alpha], dim=-1) # Piecewise: sqrt branch below 0.5, log branch above. - # Clamp inside the log branch so negative / out-of-range values don't blow up; + # Clamp the log branch at the 0.5 branch point (not above it) so the + # unselected lane stays finite in exp() without altering selected values; # values above 1.0 are allowed and extrapolate naturally. low = (t ** 2) / 3.0 - high = (torch.exp((t.clamp(min=_HLG_C) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0 + high = (torch.exp((t.clamp(min=0.5) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0 return torch.where(t <= 0.5, low, high) From ec0e8b3447d5aa5a91a5a846b7fd94c88318fef7 Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:38:50 +0300 Subject: [PATCH 25/37] fix(image): support single-channel images in Save Image (Advanced) (#14761) Signed-off-by: bigcat88 Co-authored-by: Alexis Rolland --- comfy_extras/nodes_images.py | 32 +++++++++++++++++++++----------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index 4d7b37200..7011d9c13 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -844,15 +844,18 @@ class ImageMergeTileList(IO.ComfyNode): # Format specifications # --------------------------------------------------------------------------- -# Maps (file_format, bit_depth, has_alpha) -> (numpy dtype scale, av pixel format, -# stream pix_fmt). Keeps the encode path declarative instead of branchy. +# Maps (file_format, bit_depth, num_channels) -> (quantization scale, numpy dtype, +# av frame pix_fmt, stream pix_fmt). Keeps the encode path declarative instead of branchy. _FORMAT_SPECS = { - ("png", "8-bit", False): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"}, - ("png", "8-bit", True): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"}, - ("png", "16-bit", False): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"}, - ("png", "16-bit", True): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"}, - ("exr", "32-bit float", False): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"}, - ("exr", "32-bit float", True): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"}, + ("png", "8-bit", 1): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "gray", "stream_fmt": "gray"}, + ("png", "8-bit", 3): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"}, + ("png", "8-bit", 4): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"}, + ("png", "16-bit", 1): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "gray16le", "stream_fmt": "gray16be"}, + ("png", "16-bit", 3): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"}, + ("png", "16-bit", 4): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"}, + ("exr", "32-bit float", 1): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "grayf32le", "stream_fmt": "grayf32le"}, + ("exr", "32-bit float", 3): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"}, + ("exr", "32-bit float", 4): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"}, } @@ -1088,7 +1091,8 @@ def _encode_image( bit_depth: str, colorspace: str, ) -> bytes: - """Encode a single HxWxC tensor to PNG or EXR bytes in memory. + """Encode a single HxWxC (or channel-less HxW grayscale) tensor to PNG or + EXR bytes in memory. Grayscale is written as single-channel PNG / Y-only EXR. For EXR the input is interpreted according to `colorspace` and converted to scene-linear (EXR's convention) before writing: @@ -1102,10 +1106,16 @@ def _encode_image( For PNG, colorspace selection does not modify pixels — PNG is delivered sRGB-encoded and there is no PNG path for wide-gamut HDR in this node. """ + if img_tensor.ndim == 2: + img_tensor = img_tensor.unsqueeze(-1) # Some nodes emit grayscale as (H, W) with no channel dim, mask-style. height, width, num_channels = img_tensor.shape - has_alpha = num_channels == 4 - spec = _FORMAT_SPECS[(file_format, bit_depth, has_alpha)] + spec = _FORMAT_SPECS.get((file_format, bit_depth, num_channels)) + if spec is None: + raise ValueError( + f"No {file_format}/{bit_depth} encoder for {num_channels}-channel images: " + "supported channel counts are 1 (grayscale), 3 (RGB) and 4 (RGBA)." + ) if spec["dtype"] == np.float32: # EXR path: preserve full range, no clamp. From 5697b970173bc0c16a05c30d509d0911f2b84822 Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:18:09 +0300 Subject: [PATCH 26/37] [Partner Nodes] chore(Google): reroute Gemini Image preview models to release versions (#14917) Signed-off-by: bigcat88 --- comfy_api_nodes/nodes_gemini.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/comfy_api_nodes/nodes_gemini.py b/comfy_api_nodes/nodes_gemini.py index aa992802d..a8eb0a797 100644 --- a/comfy_api_nodes/nodes_gemini.py +++ b/comfy_api_nodes/nodes_gemini.py @@ -1133,7 +1133,9 @@ class GeminiImage2(IO.ComfyNode): ) -> IO.NodeOutput: validate_string(prompt, strip_whitespace=True, min_length=1) if model == "Nano Banana 2 (Gemini 3.1 Flash Image)": - model = "gemini-3.1-flash-image-preview" + model = "gemini-3.1-flash-image" + elif model == "gemini-3-pro-image-preview": + model = "gemini-3-pro-image" parts: list[GeminiPart] = [GeminiPart(text=prompt)] if images is not None: @@ -1507,7 +1509,7 @@ class GeminiNanoBanana2V2(IO.ComfyNode): validate_string(prompt, strip_whitespace=True, min_length=1) model_choice = model["model"] if model_choice == "Nano Banana 2 (Gemini 3.1 Flash Image)": - model_id = "gemini-3.1-flash-image-preview" + model_id = "gemini-3.1-flash-image" elif model_choice == "Nano Banana 2 Lite": model_id = "gemini-3.1-flash-lite-image" else: From 5bb831a3f565dbcd517fc157284ce69af528ec2e Mon Sep 17 00:00:00 2001 From: "Daxiong (Lin)" Date: Mon, 13 Jul 2026 22:13:23 +0800 Subject: [PATCH 27/37] chore: update embedded docs to v0.5.8 (#14920) --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 790ef4940..b27de8987 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ comfyui-frontend-package==1.45.20 comfyui-workflow-templates==0.11.6 -comfyui-embedded-docs==0.5.7 +comfyui-embedded-docs==0.5.8 torch torchsde torchvision From 5658a68a875e4c1210d72c71d61eca95740adaf0 Mon Sep 17 00:00:00 2001 From: Comfy Org PR Bot Date: Tue, 14 Jul 2026 02:20:58 +0900 Subject: [PATCH 28/37] chore(openapi): sync shared API contract from cloud@bcb8f5f (#14815) Co-authored-by: mattmillerai <7741082+mattmillerai@users.noreply.github.com> Co-authored-by: Alexis Rolland Co-authored-by: Matt Miller --- openapi.yaml | 45 ++++++++++++++++++++++++--------------------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/openapi.yaml b/openapi.yaml index c09b1eeac..e00643bad 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -7,18 +7,18 @@ components: description: Timestamp when the asset was created format: date-time type: string + display_name: + description: Display name of the asset. Mirrors name for backwards compatibility. + nullable: true + type: string + file_path: + description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors") + nullable: true + type: string hash: description: Blake3 hash of the asset content. pattern: ^blake3:[a-f0-9]{64}$ type: string - loader_path: - description: The value a loader consumes to load this asset. Null when no loader can resolve the file. - nullable: true - type: string - display_name: - description: Human-facing label for the asset. Not unique. - nullable: true - type: string id: description: Unique identifier for the asset format: uuid @@ -144,6 +144,14 @@ components: AssetUpdated: description: Response returned when an existing asset is successfully updated. properties: + display_name: + description: Display name of the asset. Mirrors name for backwards compatibility. + nullable: true + type: string + file_path: + description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors") + nullable: true + type: string hash: description: Blake3 hash of the asset content. pattern: ^blake3:[a-f0-9]{64}$ @@ -775,14 +783,6 @@ components: ModelFolder: description: Represents a folder containing models properties: - extensions: - description: The folder's registered file-extension allowlist. An empty array means the folder accepts any extension (match-all). - example: - - .ckpt - - .safetensors - items: - type: string - type: array folders: description: List of paths where models of this type are stored example: @@ -1644,7 +1644,7 @@ paths: format: uuid type: string tags: - description: JSON-encoded array of tag strings. For new byte uploads, include exactly one destination role (`input`, `output`, or `models`); `models` uploads also require exactly one `model_type:` tag. Extra tags are stored as labels and do not create path components. + description: JSON-encoded array of freeform tag strings, e.g. '["models","checkpoint"]'. Common types include "models", "input", "output", and "temp", but any tag can be used in any order. type: string user_metadata: description: Custom JSON metadata as a string @@ -1829,7 +1829,7 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/Asset' + $ref: '#/components/schemas/AssetUpdated' description: Asset updated successfully "400": content: @@ -2470,9 +2470,6 @@ paths: supports_preview_metadata: description: Whether the server supports preview metadata type: boolean - supports_model_type_tags: - description: Whether the server supports namespaced model type asset tags - type: boolean type: object description: Success headers: @@ -3300,6 +3297,12 @@ paths: schema: $ref: '#/components/schemas/ErrorResponse' description: Invalid request parameters + "401": + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' + description: Unauthorized - Authentication required "500": content: application/json: From da2608926eaf68fd532bba4e1ace3402c5d21399 Mon Sep 17 00:00:00 2001 From: "Daxiong (Lin)" Date: Tue, 14 Jul 2026 03:15:34 +0800 Subject: [PATCH 29/37] Update workflow templates to v0.11.9 (#14924) --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index b27de8987..e7e7ba747 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ comfyui-frontend-package==1.45.20 -comfyui-workflow-templates==0.11.6 +comfyui-workflow-templates==0.11.9 comfyui-embedded-docs==0.5.8 torch torchsde From c35a622acdabf99f60bba737f07977fef1ef97f4 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Mon, 13 Jul 2026 12:52:28 -0700 Subject: [PATCH 30/37] Fix hidream o1 regression. (#14923) --- comfy/ldm/hidream_o1/attention.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/comfy/ldm/hidream_o1/attention.py b/comfy/ldm/hidream_o1/attention.py index 1b68f1771..afb2be9b8 100644 --- a/comfy/ldm/hidream_o1/attention.py +++ b/comfy/ldm/hidream_o1/attention.py @@ -15,24 +15,24 @@ def make_two_pass_attention(ar_len: int, transformer_options=None): The AR pass goes through SDPA directand bypasses wrappers, it is only ~1% of T at typical edit sizes. """ - def two_pass_attention(q, k, v, heads, **kwargs): + def two_pass_attention(q, k, v, heads, enable_gqa=False, **kwargs): B, H, T, D = q.shape if T < k.shape[2]: # KV-cache hot path: Q is shorter than K/V (cached AR prefix is in K/V only), all fresh Q positions are in the gen region, single full-attention call - out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options) + out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa) elif ar_len >= T: - out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True) + out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa) elif ar_len <= 0: - out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options) + out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa) else: out_ar = comfy.ops.scaled_dot_product_attention( q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len], - attn_mask=None, dropout_p=0.0, is_causal=True, + attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa, ) out_gen = optimized_attention( q[:, :, ar_len:], k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, - transformer_options=transformer_options, + transformer_options=transformer_options, enable_gqa=enable_gqa, ) out = torch.cat([out_ar, out_gen], dim=2) From 80acfcf0fe6ea006f0d6176be8301ade1c0c4836 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Mon, 13 Jul 2026 13:03:29 -0700 Subject: [PATCH 31/37] More optimized int8 and int4 on turing. (#14927) --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index e7e7ba747..e1458ca34 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,7 +22,7 @@ alembic SQLAlchemy>=2.0.0 filelock av>=16.0.0 -comfy-kitchen==0.2.18 +comfy-kitchen==0.2.19 comfy-aimdo==0.4.10 requests simpleeval>=1.0.0 From 0aecac867d7840b56ad790aa76c5e76e33c74c3d Mon Sep 17 00:00:00 2001 From: Terry Jia Date: Mon, 13 Jul 2026 22:07:23 -0400 Subject: [PATCH 32/37] Fix 3D advanced nodes crashing in API mode (no UI) (#14930) --- comfy_extras/nodes_load_3d.py | 12 ++++++++---- comfy_extras/nodes_save_3d.py | 3 ++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/comfy_extras/nodes_load_3d.py b/comfy_extras/nodes_load_3d.py index a9df557c2..106b01f9d 100644 --- a/comfy_extras/nodes_load_3d.py +++ b/comfy_extras/nodes_load_3d.py @@ -174,8 +174,9 @@ class Preview3DAdvanced(IO.ComfyNode): filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}" model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename)) + viewport_state = viewport_state if isinstance(viewport_state, dict) else {} camera_info_input = kwargs.get("camera_info", None) - camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info'] + camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info') model_3d_info_input = kwargs.get("model_3d_info", None) model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', []) return IO.NodeOutput( @@ -243,8 +244,9 @@ class PreviewGaussianSplat(IO.ComfyNode): filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}" model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename)) + viewport_state = viewport_state if isinstance(viewport_state, dict) else {} camera_info_input = kwargs.get("camera_info", None) - camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info'] + camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info') model_3d_info_input = kwargs.get("model_3d_info", None) model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', []) return IO.NodeOutput( @@ -303,8 +305,9 @@ class PreviewPointCloud(IO.ComfyNode): filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}" model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename)) + viewport_state = viewport_state if isinstance(viewport_state, dict) else {} camera_info_input = kwargs.get("camera_info", None) - camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info'] + camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info') model_3d_info_input = kwargs.get("model_3d_info", None) model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', []) return IO.NodeOutput( @@ -375,8 +378,9 @@ class Load3DAdvanced(IO.ComfyNode): file_3d = None if model_file and model_file != "none": file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file)) + viewport_state = viewport_state if isinstance(viewport_state, dict) else {} model_3d_info = viewport_state.get('model_3d_info', []) - return IO.NodeOutput(file_3d, model_3d_info, viewport_state['camera_info'], width, height) + return IO.NodeOutput(file_3d, model_3d_info, viewport_state.get('camera_info'), width, height) class Load3DExtension(ComfyExtension): diff --git a/comfy_extras/nodes_save_3d.py b/comfy_extras/nodes_save_3d.py index 7c524caa1..e9fd07326 100644 --- a/comfy_extras/nodes_save_3d.py +++ b/comfy_extras/nodes_save_3d.py @@ -418,8 +418,9 @@ def _save_file3d_to_output(model_3d: Types.File3D, filename_prefix: str) -> str: def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput: model_file = _save_file3d_to_output(model_3d, filename_prefix) + viewport_state = viewport_state if isinstance(viewport_state, dict) else {} camera_info_input = kwargs.get("camera_info", None) - camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info'] + camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info') model_3d_info_input = kwargs.get("model_3d_info", None) model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', []) return IO.NodeOutput( From dff0b18fff04fda6558d0c1dd2ad1dca43fc18bc Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Tue, 14 Jul 2026 13:13:08 -0700 Subject: [PATCH 33/37] Fix int8 performance regression on 16xx series. (#14941) --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index e1458ca34..eb40caa6b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,7 +22,7 @@ alembic SQLAlchemy>=2.0.0 filelock av>=16.0.0 -comfy-kitchen==0.2.19 +comfy-kitchen==0.2.20 comfy-aimdo==0.4.10 requests simpleeval>=1.0.0 From 26537080cb1da5a6efb05bfc5a4792569c845913 Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Tue, 14 Jul 2026 23:21:20 +0300 Subject: [PATCH 34/37] [Partner Nodes] feat(sync.so): add support for "sync-3" model (#14928) Signed-off-by: bigcat88 --- comfy_api_nodes/apis/sync_so.py | 49 ++++ comfy_api_nodes/nodes_sync_so.py | 391 +++++++++++++++++++++++++++++++ 2 files changed, 440 insertions(+) create mode 100644 comfy_api_nodes/apis/sync_so.py create mode 100644 comfy_api_nodes/nodes_sync_so.py diff --git a/comfy_api_nodes/apis/sync_so.py b/comfy_api_nodes/apis/sync_so.py new file mode 100644 index 000000000..af9419580 --- /dev/null +++ b/comfy_api_nodes/apis/sync_so.py @@ -0,0 +1,49 @@ +from pydantic import BaseModel, Field + + +class SyncInputItem(BaseModel): + type: str = Field(..., description="Input kind: 'video', 'image' or 'audio'.") + url: str = Field(...) + + +class SyncActiveSpeakerDetection(BaseModel): + auto_detect: bool | None = Field( + None, description="Detect the active speaker automatically. Video input only; rejected for images." + ) + frame_number: int | None = Field( + None, description="Frame used for manual speaker selection. Must be 0 for image inputs." + ) + coordinates: list[int] | None = Field( + None, description="Pixel [x, y] of the speaker's face in the frame selected by frame_number." + ) + + +class SyncGenerationOptions(BaseModel): + sync_mode: str | None = Field( + None, + description="How to resolve an audio/video duration mismatch: " + "cut_off, bounce, loop, silence or remap. Ignored for image inputs.", + ) + i2v_prompt: str | None = Field( + None, description="Motion prompt for image-to-video generation. Image input only." + ) + active_speaker_detection: SyncActiveSpeakerDetection | None = Field(None) + + +class SyncGenerationRequest(BaseModel): + model: str = Field(..., description="Generation model, e.g. 'sync-3'.") + input: list[SyncInputItem] = Field( + ..., description="Exactly one visual input (video or image) plus one audio input." + ) + options: SyncGenerationOptions | None = Field(None) + + +class SyncGeneration(BaseModel): + """Subset of the Generation object returned by POST /v2/generate and GET /v2/generate/{id}.""" + + id: str = Field(...) + status: str = Field(..., description="PENDING | PROCESSING | COMPLETED | FAILED | REJECTED") + outputUrl: str | None = Field(None) + outputDuration: float | None = Field(None) + error: str | None = Field(None, description="Human-readable failure message.") + errorCode: str | None = Field(None, description="Stable machine-readable code from the GET /v2/errors catalog.") diff --git a/comfy_api_nodes/nodes_sync_so.py b/comfy_api_nodes/nodes_sync_so.py new file mode 100644 index 000000000..27382b399 --- /dev/null +++ b/comfy_api_nodes/nodes_sync_so.py @@ -0,0 +1,391 @@ +from typing_extensions import override + +from comfy_api.latest import IO, ComfyExtension, Input +from comfy_api_nodes.apis.sync_so import ( + SyncActiveSpeakerDetection, + SyncGeneration, + SyncGenerationOptions, + SyncGenerationRequest, + SyncInputItem, +) +from comfy_api_nodes.util import ( + ApiEndpoint, + download_url_to_video_output, + downscale_image_tensor, + downscale_image_tensor_by_max_side, + get_image_dimensions, + get_number_of_images, + poll_op, + sync_op, + upload_audio_to_comfyapi, + upload_image_to_comfyapi, + upload_video_to_comfyapi, + validate_audio_duration, +) + + +class SyncLipSyncNode(IO.ComfyNode): + @classmethod + def define_schema(cls) -> IO.Schema: + return IO.Schema( + node_id="SyncLipSyncNode", + display_name="sync.so Lip Sync", + category="partner/video/sync.so", + description=( + "Re-sync mouth movement in a video to new speech audio using sync.so. " + "Handles close-ups, profiles and obstructions automatically while preserving " + "the speaker's expression. Cost scales with output duration." + ), + inputs=[ + IO.Video.Input( + "video", + tooltip="Footage of the speaker to re-sync. Up to 4K (4096x2160); " + "a constant frame rate of 24/25/30 fps works best.", + ), + IO.Audio.Input( + "audio", + tooltip="Speech audio to sync the mouth to.", + ), + IO.Int.Input( + "seed", + default=42, + min=0, + max=2147483647, + control_after_generate=True, + tooltip="Seed controls whether the node should re-run; " + "results are non-deterministic regardless of seed.", + ), + IO.DynamicCombo.Input( + "model", + options=[ + IO.DynamicCombo.Option( + "sync-3", + [ + IO.Combo.Input( + "sync_mode", + options=["bounce", "cut_off", "loop", "silence", "remap"], + default="bounce", + tooltip=( + "How to handle a duration mismatch between video and audio; " + "this also sets the output length. " + "bounce: video plays forward then backward until the audio ends " + "(output = audio length). " + "loop: video restarts until the audio ends (output = audio length). " + "remap: video is time-stretched to match the audio (output = audio length). " + "cut_off: the longer track is trimmed (output = shorter length). " + "silence: nothing is trimmed; the shorter track is padded " + "(output = longer length)." + ), + ), + IO.Combo.Input( + "speaker_selection", + options=["default", "auto-detect", "coordinates"], + default="default", + tooltip=( + "Which face to lipsync when several people are visible. " + "default: let the model decide. " + "auto-detect: detect and follow the active speaker. " + "coordinates: target the face at pixel (speaker_x, speaker_y) " + "in the frame chosen by speaker_frame." + ), + ), + IO.Int.Input( + "speaker_frame", + default=0, + min=0, + max=1_000_000, + advanced=True, + tooltip="Video frame used to locate the speaker. " + "Only used when speaker_selection is 'coordinates'.", + ), + IO.Int.Input( + "speaker_x", + default=0, + min=0, + max=4096, + advanced=True, + tooltip="X pixel coordinate of the speaker's face. " + "Only used when speaker_selection is 'coordinates'.", + ), + IO.Int.Input( + "speaker_y", + default=0, + min=0, + max=4096, + advanced=True, + tooltip="Y pixel coordinate of the speaker's face. " + "Only used when speaker_selection is 'coordinates'.", + ), + ], + ) + ], + tooltip="sync.so generation model.", + ), + ], + outputs=[IO.Video.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( + expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""", + ), + ) + + @classmethod + async def execute( + cls, + video: Input.Video, + audio: Input.Audio, + seed: int, + model: dict, + ) -> IO.NodeOutput: + try: + width, height = video.get_dimensions() + except Exception: + width = height = None + if width and height and (max(width, height) > 4096 or width * height > 4096 * 2160): + raise ValueError( + f"sync.so rejects videos above 4K (4096x2160); got {width}x{height}. Downscale the video first." + ) + validate_audio_duration(audio, max_duration=600) + + if model["speaker_selection"] == "auto-detect": + speaker_detection = SyncActiveSpeakerDetection(auto_detect=True) + elif model["speaker_selection"] == "coordinates": + speaker_detection = SyncActiveSpeakerDetection( + frame_number=model["speaker_frame"], + coordinates=[model["speaker_x"], model["speaker_y"]], + ) + else: + speaker_detection = None + + video_url = await upload_video_to_comfyapi(cls, video, max_duration=600) + audio_url = await upload_audio_to_comfyapi(cls, audio) + + generation = await sync_op( + cls, + ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"), + response_model=SyncGeneration, + data=SyncGenerationRequest( + model=model["model"], + input=[ + SyncInputItem(type="video", url=video_url), + SyncInputItem(type="audio", url=audio_url), + ], + options=SyncGenerationOptions( + sync_mode=model["sync_mode"], + active_speaker_detection=speaker_detection, + ), + ), + ) + generation = await poll_op( + cls, + ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"), + response_model=SyncGeneration, + status_extractor=lambda g: g.status, + completed_statuses=["COMPLETED", "FAILED", "REJECTED"], + failed_statuses=[], + queued_statuses=["PENDING"], + poll_interval=10.0, + ) + if generation.status != "COMPLETED": + code = f" [{generation.errorCode}]" if generation.errorCode else "" + raise ValueError( + f"sync.so generation {generation.status.lower()}{code}: " + f"{generation.error or 'no error details provided'}" + ) + if not generation.outputUrl: + raise ValueError("sync.so generation completed but no output URL was returned.") + return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl)) + + +class SyncTalkingImageNode(IO.ComfyNode): + @classmethod + def define_schema(cls) -> IO.Schema: + return IO.Schema( + node_id="SyncTalkingImageNode", + display_name="sync.so Talking Image", + category="partner/video/sync.so", + description=( + "Animate a still portrait into a talking video driven by speech audio, " + "using sync.so's sync-3 model. The output duration matches the audio. " + "Cost scales with output duration." + ), + inputs=[ + IO.Image.Input( + "image", + tooltip="A single image with a clearly visible face, up to 4K (4096x2160).", + ), + IO.Audio.Input( + "audio", + tooltip="Speech audio driving the talking video; the output duration matches it. " + "Chain any TTS node here to drive the animation from text.", + ), + IO.String.Input( + "prompt", + multiline=True, + default="", + tooltip="Optional guidance for how the portrait comes to life, e.g. " + "'make the subject smile and look at the camera'. " + "Leave empty for natural talking motion.", + ), + IO.Int.Input( + "seed", + default=0, + min=0, + max=2147483647, + control_after_generate=True, + tooltip="Seed controls whether the node should re-run; " + "results are non-deterministic regardless of seed.", + ), + IO.DynamicCombo.Input( + "model", + options=[ + IO.DynamicCombo.Option( + "sync-3", + [ + IO.Combo.Input( + "speaker_selection", + options=["default", "coordinates"], + default="default", + tooltip=( + "Which face to animate when several people are visible. " + "default: let the model decide. " + "coordinates: target the face at pixel (speaker_x, speaker_y) " + "in the image. Auto-detection is not supported for images." + ), + ), + IO.Int.Input( + "speaker_x", + default=0, + min=0, + max=4096, + advanced=True, + tooltip="X pixel coordinate of the speaker's face. " + "Only used when speaker_selection is 'coordinates'.", + ), + IO.Int.Input( + "speaker_y", + default=0, + min=0, + max=4096, + advanced=True, + tooltip="Y pixel coordinate of the speaker's face. " + "Only used when speaker_selection is 'coordinates'.", + ), + IO.Boolean.Input( + "auto_downscale", + default=True, + advanced=True, + tooltip="Automatically downscale the image if it exceeds the 4K " + "(4096x2160) input limit; speaker coordinates are scaled to match. " + "When disabled, an oversized image raises an error instead.", + ), + ], + ) + ], + tooltip="sync.so generation model. Image input is exclusive to sync-3.", + ), + ], + outputs=[IO.Video.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( + expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""", + ), + ) + + @classmethod + async def execute( + cls, + image: Input.Image, + audio: Input.Audio, + prompt: str, + seed: int, + model: dict, + ) -> IO.NodeOutput: + if get_number_of_images(image) != 1: + raise ValueError("Exactly one image is required; got a batch. Pick one frame first.") + validate_audio_duration(audio, max_duration=600) + + height, width = get_image_dimensions(image) + speaker_x, speaker_y = model["speaker_x"], model["speaker_y"] + if max(width, height) > 4096 or width * height > 4096 * 2160: + if not model["auto_downscale"]: + raise ValueError( + f"sync.so rejects images above 4K (4096x2160); got {width}x{height}. " + "Downscale the image first or enable auto_downscale." + ) + image = downscale_image_tensor(image, total_pixels=4096 * 2160) + image = downscale_image_tensor_by_max_side(image, max_side=4096) + new_height, new_width = get_image_dimensions(image) + # speaker coordinates are given in the original image's pixel space + speaker_x = min(new_width - 1, round(speaker_x * new_width / width)) + speaker_y = min(new_height - 1, round(speaker_y * new_height / height)) + + if model["speaker_selection"] == "coordinates": + speaker_detection = SyncActiveSpeakerDetection( + frame_number=0, # images have a single frame; auto_detect is rejected by the API + coordinates=[speaker_x, speaker_y], + ) + else: + speaker_detection = None + + image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None) + audio_url = await upload_audio_to_comfyapi(cls, audio) + + generation = await sync_op( + cls, + ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"), + response_model=SyncGeneration, + data=SyncGenerationRequest( + model=model["model"], + input=[ + SyncInputItem(type="image", url=image_url), + SyncInputItem(type="audio", url=audio_url), + ], + options=SyncGenerationOptions( + i2v_prompt=prompt.strip() or None, + active_speaker_detection=speaker_detection, + ), + ), + ) + generation = await poll_op( + cls, + ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"), + response_model=SyncGeneration, + status_extractor=lambda g: g.status, + completed_statuses=["COMPLETED", "FAILED", "REJECTED"], + failed_statuses=[], + queued_statuses=["PENDING"], + poll_interval=10.0, + ) + if generation.status != "COMPLETED": + code = f" [{generation.errorCode}]" if generation.errorCode else "" + raise ValueError( + f"sync.so generation {generation.status.lower()}{code}: " + f"{generation.error or 'no error details provided'}" + ) + if not generation.outputUrl: + raise ValueError("sync.so generation completed but no output URL was returned.") + return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl)) + + +class SyncExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[IO.ComfyNode]]: + return [ + SyncLipSyncNode, + SyncTalkingImageNode, + ] + + +async def comfy_entrypoint() -> SyncExtension: + return SyncExtension() From 1701cce8dc105e47bc1e6f5a7790c495bbcf331c Mon Sep 17 00:00:00 2001 From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com> Date: Wed, 15 Jul 2026 00:29:01 +0300 Subject: [PATCH 35/37] Fix cached outputs missing from job results when prompt has no client_id (#14939) When a prompt is submitted without client_id and its output nodes are served from cache, _send_cached_ui returned early before recording the cached UI outputs, so /api/jobs/{job_id} (and /history) reported success with empty outputs. Record the outputs before the client_id check. --- execution.py | 4 ++-- tests/execution/test_execution.py | 24 ++++++++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/execution.py b/execution.py index 19b8cdd68..387772629 100644 --- a/execution.py +++ b/execution.py @@ -426,12 +426,12 @@ def _is_intermediate_output(dynprompt, node_id): def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs): + if cached.ui is not None: + ui_outputs[node_id] = cached.ui if server.client_id is None: return cached_ui = cached.ui or {} server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id) - if cached.ui is not None: - ui_outputs[node_id] = cached.ui async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs): unique_id = current_item diff --git a/tests/execution/test_execution.py b/tests/execution/test_execution.py index 15e2304fc..c914d2feb 100644 --- a/tests/execution/test_execution.py +++ b/tests/execution/test_execution.py @@ -818,6 +818,30 @@ class TestExecution: except urllib.error.HTTPError: pass # Expected behavior + def test_cached_outputs_in_job_without_client_id(self, client: ComfyClient, builder: GraphBuilder): + g = builder + image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1) + output = g.node("SaveImage", images=image.out(0)) + + # Prime the cache with a normal run. + client.run(g) + + # Resubmit anonymously (no client_id) so output nodes are cache hits with no websocket client. + data = json.dumps({"prompt": g.finalize()}).encode('utf-8') + req = urllib.request.Request(f"http://{client.server_address}/prompt", data=data) + prompt_id = json.loads(urllib.request.urlopen(req).read())['prompt_id'] + + for _ in range(100): + job = client.get_job(prompt_id) + if job is not None and job['status'] not in ('pending', 'in_progress'): + break + time.sleep(0.1) + else: + raise AssertionError("Prompt did not complete in time") + + assert job['status'] == 'completed' + assert output.id in job['outputs'], "Cached outputs must appear in job outputs without a client_id" + def _create_history_item(self, client, builder): g = GraphBuilder(prefix="offset_test") input_node = g.node( From 3cd13eb4245a111bda4a46b9129cef19c8cac1d5 Mon Sep 17 00:00:00 2001 From: Comfy Org PR Bot Date: Wed, 15 Jul 2026 14:29:24 +0900 Subject: [PATCH 36/37] Bump comfyui-frontend-package to 1.45.21 (#14944) --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index eb40caa6b..e7d301576 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -comfyui-frontend-package==1.45.20 +comfyui-frontend-package==1.45.21 comfyui-workflow-templates==0.11.9 comfyui-embedded-docs==0.5.8 torch From 700821e1364eaab0e8f21c538a2131719fec57bf Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 15 Jul 2026 01:43:14 -0400 Subject: [PATCH 37/37] ComfyUI v0.28.0 --- comfyui_version.py | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/comfyui_version.py b/comfyui_version.py index 8e9967f1b..dcc0fee96 100644 --- a/comfyui_version.py +++ b/comfyui_version.py @@ -1,3 +1,3 @@ # This file is automatically generated by the build process when version is # updated in pyproject.toml. -__version__ = "0.27.0" +__version__ = "0.28.0" diff --git a/pyproject.toml b/pyproject.toml index 8c17e410e..73de2990f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "ComfyUI" -version = "0.27.0" +version = "0.28.0" readme = "README.md" license = { file = "LICENSE" } requires-python = ">=3.10"