Merge branch 'master' into worksplit-multigpu

comfy/patcher_extension.py

@@ -50,6 +50,7 @@ def get_all_callbacks(call_type: str, transformer_options: dict, is_model_option
 
 class WrappersMP:
     OUTER_SAMPLE = "outer_sample"
+    PREPARE_SAMPLING = "prepare_sampling"
     SAMPLER_SAMPLE = "sampler_sample"
     CALC_COND_BATCH = "calc_cond_batch"
     APPLY_MODEL = "apply_model"

comfy/sampler_helpers.py

@@ -148,6 +148,13 @@ def preprocess_multigpu_conds(conds: dict[str, list[dict[str]]], model: ModelPat
     # potentially handle gligen - since not widely used, ignored for now
 
 def prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None):
+    executor = comfy.patcher_extension.WrapperExecutor.new_executor(
+        _prepare_sampling,
+        comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.PREPARE_SAMPLING, model_options, is_model_options=True)
+    )
+    return executor.execute(model, noise_shape, conds, model_options=model_options)
+
+def _prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None):
     model.match_multigpu_clones()
     preprocess_multigpu_conds(conds, model, model_options)
     models, inference_memory = get_additional_models(conds, model.model_dtype())

comfy_extras/nodes_hunyuan3d.py

@@ -209,6 +209,196 @@ def voxel_to_mesh(voxels, threshold=0.5, device=None):
     vertices = torch.fliplr(vertices)
     return vertices, faces
 
+def voxel_to_mesh_surfnet(voxels, threshold=0.5, device=None):
+    if device is None:
+        device = torch.device("cpu")
+    voxels = voxels.to(device)
+
+    D, H, W = voxels.shape
+
+    padded = torch.nn.functional.pad(voxels, (1, 1, 1, 1, 1, 1), 'constant', 0)
+    z, y, x = torch.meshgrid(
+        torch.arange(D, device=device),
+        torch.arange(H, device=device),
+        torch.arange(W, device=device),
+        indexing='ij'
+    )
+    cell_positions = torch.stack([z.flatten(), y.flatten(), x.flatten()], dim=1)
+
+    corner_offsets = torch.tensor([
+        [0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0],
+        [0, 0, 1], [1, 0, 1], [0, 1, 1], [1, 1, 1]
+    ], device=device)
+
+    corner_values = torch.zeros((cell_positions.shape[0], 8), device=device)
+    for c, (dz, dy, dx) in enumerate(corner_offsets):
+        corner_values[:, c] = padded[
+            cell_positions[:, 0] + dz,
+            cell_positions[:, 1] + dy,
+            cell_positions[:, 2] + dx
+        ]
+
+    corner_signs = corner_values > threshold
+    has_inside = torch.any(corner_signs, dim=1)
+    has_outside = torch.any(~corner_signs, dim=1)
+    contains_surface = has_inside & has_outside
+
+    active_cells = cell_positions[contains_surface]
+    active_signs = corner_signs[contains_surface]
+    active_values = corner_values[contains_surface]
+
+    if active_cells.shape[0] == 0:
+        return torch.zeros((0, 3), device=device), torch.zeros((0, 3), dtype=torch.long, device=device)
+
+    edges = torch.tensor([
+        [0, 1], [0, 2], [0, 4], [1, 3],
+        [1, 5], [2, 3], [2, 6], [3, 7],
+        [4, 5], [4, 6], [5, 7], [6, 7]
+    ], device=device)
+
+    cell_vertices = {}
+    progress = comfy.utils.ProgressBar(100)
+
+    for edge_idx, (e1, e2) in enumerate(edges):
+        progress.update(1)
+        crossing = active_signs[:, e1] != active_signs[:, e2]
+        if not crossing.any():
+            continue
+
+        cell_indices = torch.nonzero(crossing, as_tuple=True)[0]
+
+        v1 = active_values[cell_indices, e1]
+        v2 = active_values[cell_indices, e2]
+
+        t = torch.zeros_like(v1, device=device)
+        denom = v2 - v1
+        valid = denom != 0
+        t[valid] = (threshold - v1[valid]) / denom[valid]
+        t[~valid] = 0.5
+
+        p1 = corner_offsets[e1].float()
+        p2 = corner_offsets[e2].float()
+
+        intersection = p1.unsqueeze(0) + t.unsqueeze(1) * (p2.unsqueeze(0) - p1.unsqueeze(0))
+
+        for i, point in zip(cell_indices.tolist(), intersection):
+            if i not in cell_vertices:
+                cell_vertices[i] = []
+            cell_vertices[i].append(point)
+
+    # Calculate the final vertices as the average of intersection points for each cell
+    vertices = []
+    vertex_lookup = {}
+
+    vert_progress_mod = round(len(cell_vertices)/50)
+
+    for i, points in cell_vertices.items():
+        if not i % vert_progress_mod:
+            progress.update(1)
+
+        if points:
+            vertex = torch.stack(points).mean(dim=0)
+            vertex = vertex + active_cells[i].float()
+            vertex_lookup[tuple(active_cells[i].tolist())] = len(vertices)
+            vertices.append(vertex)
+
+    if not vertices:
+        return torch.zeros((0, 3), device=device), torch.zeros((0, 3), dtype=torch.long, device=device)
+
+    final_vertices = torch.stack(vertices)
+
+    inside_corners_mask = active_signs
+    outside_corners_mask = ~active_signs
+
+    inside_counts = inside_corners_mask.sum(dim=1, keepdim=True).float()
+    outside_counts = outside_corners_mask.sum(dim=1, keepdim=True).float()
+
+    inside_pos = torch.zeros((active_cells.shape[0], 3), device=device)
+    outside_pos = torch.zeros((active_cells.shape[0], 3), device=device)
+
+    for i in range(8):
+        mask_inside = inside_corners_mask[:, i].unsqueeze(1)
+        mask_outside = outside_corners_mask[:, i].unsqueeze(1)
+        inside_pos += corner_offsets[i].float().unsqueeze(0) * mask_inside
+        outside_pos += corner_offsets[i].float().unsqueeze(0) * mask_outside
+
+    inside_pos /= inside_counts
+    outside_pos /= outside_counts
+    gradients = inside_pos - outside_pos
+
+    pos_dirs = torch.tensor([
+        [1, 0, 0],
+        [0, 1, 0],
+        [0, 0, 1]
+    ], device=device)
+
+    cross_products = [
+        torch.linalg.cross(pos_dirs[i].float(), pos_dirs[j].float())
+        for i in range(3) for j in range(i+1, 3)
+    ]
+
+    faces = []
+    all_keys = set(vertex_lookup.keys())
+
+    face_progress_mod = round(len(active_cells)/38*3)
+
+    for pair_idx, (i, j) in enumerate([(0,1), (0,2), (1,2)]):
+        dir_i = pos_dirs[i]
+        dir_j = pos_dirs[j]
+        cross_product = cross_products[pair_idx]
+
+        ni_positions = active_cells + dir_i
+        nj_positions = active_cells + dir_j
+        diag_positions = active_cells + dir_i + dir_j
+
+        alignments = torch.matmul(gradients, cross_product)
+
+        valid_quads = []
+        quad_indices = []
+
+        for idx, active_cell in enumerate(active_cells):
+            if not idx % face_progress_mod:
+                progress.update(1)
+            cell_key = tuple(active_cell.tolist())
+            ni_key = tuple(ni_positions[idx].tolist())
+            nj_key = tuple(nj_positions[idx].tolist())
+            diag_key = tuple(diag_positions[idx].tolist())
+
+            if cell_key in all_keys and ni_key in all_keys and nj_key in all_keys and diag_key in all_keys:
+                v0 = vertex_lookup[cell_key]
+                v1 = vertex_lookup[ni_key]
+                v2 = vertex_lookup[nj_key]
+                v3 = vertex_lookup[diag_key]
+
+                valid_quads.append((v0, v1, v2, v3))
+                quad_indices.append(idx)
+
+        for q_idx, (v0, v1, v2, v3) in enumerate(valid_quads):
+            cell_idx = quad_indices[q_idx]
+            if alignments[cell_idx] > 0:
+                faces.append(torch.tensor([v0, v1, v3], device=device, dtype=torch.long))
+                faces.append(torch.tensor([v0, v3, v2], device=device, dtype=torch.long))
+            else:
+                faces.append(torch.tensor([v0, v3, v1], device=device, dtype=torch.long))
+                faces.append(torch.tensor([v0, v2, v3], device=device, dtype=torch.long))
+
+    if faces:
+        faces = torch.stack(faces)
+    else:
+        faces = torch.zeros((0, 3), dtype=torch.long, device=device)
+
+    v_min = 0
+    v_max = max(D, H, W)
+
+    final_vertices = final_vertices - (v_min + v_max) / 2
+
+    scale = (v_max - v_min) / 2
+    if scale > 0:
+        final_vertices = final_vertices / scale
+
+    final_vertices = torch.fliplr(final_vertices)
+
+    return final_vertices, faces
 
 class MESH:
     def __init__(self, vertices, faces):
@@ -237,6 +427,34 @@ class VoxelToMeshBasic:
 
         return (MESH(torch.stack(vertices), torch.stack(faces)), )
 
+class VoxelToMesh:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"voxel": ("VOXEL", ),
+                             "algorithm": (["surface net", "basic"], ),
+                             "threshold": ("FLOAT", {"default": 0.6, "min": -1.0, "max": 1.0, "step": 0.01}),
+                             }}
+    RETURN_TYPES = ("MESH",)
+    FUNCTION = "decode"
+
+    CATEGORY = "3d"
+
+    def decode(self, voxel, algorithm, threshold):
+        vertices = []
+        faces = []
+
+        if algorithm == "basic":
+            mesh_function = voxel_to_mesh
+        elif algorithm == "surface net":
+            mesh_function = voxel_to_mesh_surfnet
+
+        for x in voxel.data:
+            v, f = mesh_function(x, threshold=threshold, device=None)
+            vertices.append(v)
+            faces.append(f)
+
+        return (MESH(torch.stack(vertices), torch.stack(faces)), )
+
 
 def save_glb(vertices, faces, filepath, metadata=None):
     """
@@ -411,5 +629,6 @@ NODE_CLASS_MAPPINGS = {
     "Hunyuan3Dv2ConditioningMultiView": Hunyuan3Dv2ConditioningMultiView,
     "VAEDecodeHunyuan3D": VAEDecodeHunyuan3D,
     "VoxelToMeshBasic": VoxelToMeshBasic,
+    "VoxelToMesh": VoxelToMesh,
     "SaveGLB": SaveGLB,
 }

execution.py

@@ -775,7 +775,7 @@ def validate_prompt(prompt):
                 "details": f"Node ID '#{x}'",
                 "extra_info": {}
             }
-            return (False, error, [], [])
+            return (False, error, [], {})
 
         class_type = prompt[x]['class_type']
         class_ = nodes.NODE_CLASS_MAPPINGS.get(class_type, None)
@@ -786,7 +786,7 @@ def validate_prompt(prompt):
                 "details": f"Node ID '#{x}'",
                 "extra_info": {}
             }
-            return (False, error, [], [])
+            return (False, error, [], {})
 
         if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE is True:
             outputs.add(x)
@@ -798,7 +798,7 @@ def validate_prompt(prompt):
             "details": "",
             "extra_info": {}
         }
-        return (False, error, [], [])
+        return (False, error, [], {})
 
     good_outputs = set()
     errors = []

nodes.py

@@ -1692,6 +1692,9 @@ class LoadImage:
             if 'A' in i.getbands():
                 mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
                 mask = 1. - torch.from_numpy(mask)
+            elif i.mode == 'P' and 'transparency' in i.info:
+                mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
+                mask = 1. - torch.from_numpy(mask)
             else:
                 mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
             output_images.append(image)
@@ -2127,21 +2130,25 @@ def get_module_name(module_path: str) -> str:
 
 
 def load_custom_node(module_path: str, ignore=set(), module_parent="custom_nodes") -> bool:
-    module_name = os.path.basename(module_path)
+    module_name = get_module_name(module_path)
     if os.path.isfile(module_path):
         sp = os.path.splitext(module_path)
         module_name = sp[0]
+        sys_module_name = module_name
+    elif os.path.isdir(module_path):
+        sys_module_name = module_path
+
     try:
         logging.debug("Trying to load custom node {}".format(module_path))
         if os.path.isfile(module_path):
-            module_spec = importlib.util.spec_from_file_location(module_name, module_path)
+            module_spec = importlib.util.spec_from_file_location(sys_module_name, module_path)
             module_dir = os.path.split(module_path)[0]
         else:
-            module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
+            module_spec = importlib.util.spec_from_file_location(sys_module_name, os.path.join(module_path, "__init__.py"))
             module_dir = module_path
 
         module = importlib.util.module_from_spec(module_spec)
-        sys.modules[module_name] = module
+        sys.modules[sys_module_name] = module
         module_spec.loader.exec_module(module)
 
         LOADED_MODULE_DIRS[module_name] = os.path.abspath(module_dir)

server.py

@@ -657,7 +657,13 @@ class PromptServer():
                     logging.warning("invalid prompt: {}".format(valid[1]))
                     return web.json_response({"error": valid[1], "node_errors": valid[3]}, status=400)
             else:
-                return web.json_response({"error": "no prompt", "node_errors": []}, status=400)
+                error = {
+                    "type": "no_prompt",
+                    "message": "No prompt provided",
+                    "details": "No prompt provided",
+                    "extra_info": {}
+                }
+                return web.json_response({"error": error, "node_errors": {}}, status=400)
 
         @routes.post("/queue")
         async def post_queue(request):