Fix inputs

infinigen/assets/objects/creatures/parts/ridged_fin.py

@@ -12,7 +12,7 @@
 from infinigen.assets.materials.utils.surface_utils import sample_range
 from infinigen.assets.objects.creatures.util.creature import Part, PartFactory
 from infinigen.core.nodes import node_utils
-from infinigen.core.nodes.node_wrangler import Nodes, NodeWrangler
+from infinigen.core.nodes.node_wrangler import Nodes, NodeWrangler, ng_outputs
 from infinigen.core.tagging import tag_object
 from infinigen.core.util import blender as butil
 
@@ -649,17 +649,9 @@ def make_part(self, params):
         )
         butil.set_geomod_inputs(mod, params)
 
-        id1 = next(
-            s
-            for s in mod.node_group.interface.items_tree
-            if s.in_out == "OUTPUT" and s.name == "Bump"
-        ).identifier
+        id1 = ng_outputs(mod.node_group)["Bump"].identifier
         mod[f"{id1}_attribute_name"] = "Bump"
-        id2 = next(
-            s
-            for s in mod.node_group.interface.items_tree
-            if s.in_out == "OUTPUT" and s.name == "BumpMask"
-        ).identifier
+        id2 = ng_outputs(mod.node_group)["BumpMask"].identifier
         mod[f"{id2}_attribute_name"] = "BumpMask"
 
         butil.apply_modifiers(fin, mod)

infinigen/assets/objects/creatures/util/part_util.py

@@ -15,6 +15,7 @@
 )
 from infinigen.assets.utils.extract_nodegroup_parts import extract_nodegroup_geo
 from infinigen.assets.utils.geometry import nurbs
+from infinigen.core.nodes.node_wrangler import ng_outputs
 from infinigen.core.util import blender as butil
 
 
@@ -29,11 +30,11 @@ def nodegroup_to_part(
             ng = nodegroup_func(**kwargs)
         else:
             ng = nodegroup_func()
-        geo_outputs = [
-            s
-            for s in ng.interface.items_tree
-            if s.in_out == "OUTPUT" and s.bl_socket_idname == "NodeSocketGeometry"
-        ]
+        geo_outputs = {
+            o
+            for o in ng_outputs(ng).values()
+            if o.bl_socket_idname == "NodeSocketGeometry"
+        }
         objs = {
             o.name: extract_nodegroup_geo(base_obj, ng, o.name, ng_params=params)
             for o in geo_outputs

infinigen/assets/objects/seating/sofa.py

@@ -1446,7 +1446,7 @@ def create_placeholder(self, **_):
             ng_inputs={
                 **self.params,
             },
-            apply=False,
+            apply=True,
         )
         tagging.tag_system.relabel_obj(obj)
         surface.add_material(obj, self.sofa_fabric)

infinigen/assets/utils/extract_nodegroup_parts.py

@@ -10,14 +10,14 @@
     Nodes,
     NodeWrangler,
     geometry_node_group_empty_new,
+    ng_inputs,
+    ng_outputs,
 )
 from infinigen.core.util import blender as butil
 
 
 def extract_nodegroup_geo(target_obj, nodegroup, k, ng_params=None):
-    inputs = [s.name for s in nodegroup.interface.items_tree if s.in_out == "INPUT"]
-    outputs = [s.name for s in nodegroup.interface.items_tree if s.in_out == "OUTPUT"]
-    assert k in outputs
+    assert k in ng_outputs(nodegroup)
     assert target_obj.type == "MESH"
 
     vert = butil.spawn_vert("extract_nodegroup_geo.temp")
@@ -31,7 +31,7 @@ def extract_nodegroup_geo(target_obj, nodegroup, k, ng_params=None):
     obj_inp = nw.new_node(Nodes.ObjectInfo, [target_obj])
 
     group_input_kwargs = {**ng_params}
-    if "Geometry" in inputs:
+    if "Geometry" in ng_inputs(nodegroup):
         group_input_kwargs["Geometry"] = obj_inp.outputs["Geometry"]
     group = nw.new_node(nodegroup.name, input_kwargs=group_input_kwargs)
 

infinigen/core/nodes/node_transpiler/transpiler.py

@@ -17,6 +17,7 @@
 import numpy as np
 
 from infinigen.core.nodes.node_info import OUTPUT_NODE_IDS, SINGLETON_NODES, Nodes
+from infinigen.core.nodes.node_wrangler import ng_inputs
 
 logger = logging.getLogger(__name__)
 
@@ -666,7 +667,7 @@ def create_node(node_tree, node, memo):
     # Special case: input node
     if node.bl_idname == Nodes.GroupInput:
         all_inps = []
-        for inp in node_tree.inputs:
+        for inp in ng_inputs(node_tree).values():
             repr_val, targets = (
                 represent_default_value(inp.default_value, simple=False)
                 if hasattr(inp, "default_value")

infinigen/core/nodes/node_wrangler.py

@@ -32,6 +32,14 @@ class NodeMisuseWarning(UserWarning):
     pass
 
 
+def ng_inputs(node_group):
+    return {s.name: s for s in node_group.interface.items_tree if s.in_out == "INPUT"}
+
+
+def ng_outputs(node_group):
+    return {s.name: s for s in node_group.interface.items_tree if s.in_out == "OUTPUT"}
+
+
 # This is for Blender 3.3 because of the nodetree change
 def geometry_node_group_empty_new():
     group = bpy.data.node_groups.new("Geometry Nodes", "GeometryNodeTree")
@@ -225,10 +233,7 @@ def expose_input(
 
         group_input = self.new_node(Nodes.GroupInput)  # will reuse singleton
 
-        if any(
-            s.name == name and s.in_out == "INPUT"
-            for s in self.node_group.interface.items_tree
-        ):
+        if name in ng_inputs(self.node_group):
             assert len([o for o in group_input.outputs if o.name == name]) == 1
             return group_input.outputs[name]
 

infinigen/core/surface.py

@@ -24,6 +24,8 @@
     Nodes,
     NodeWrangler,
     geometry_node_group_empty_new,
+    ng_inputs,
+    ng_outputs,
 )
 from infinigen.core.nodes.utils import infer_output_socket, isnode
 from infinigen.core.util import blender as butil
@@ -390,39 +392,31 @@ def add_geomod(
         else:
             mod.node_group = ng
 
-        outputs = [
-            s for s in mod.node_group.interface.items_tree if s.in_out == "OUTPUT"
+        non_geometries = [
+            o
+            for o in ng_outputs(mod.node_group).values()
+            if o.socket_type != "NodeSocketGeometry"
         ]
-        identifiers = [
-            outputs[i].identifier
-            for i in range(len(outputs))
-            if outputs[i].socket_type != "NodeSocketGeometry"
-        ]
-        if len(identifiers) != len(attributes):
+        if len(non_geometries) != len(attributes):
             raise Exception(
-                f"has {len(identifiers)} identifiers, but {len(attributes)} attributes. Specifically, "
-                f"{identifiers=} and {attributes=}"
+                f"has {len(non_geometries)} identifiers, but {len(attributes)} attributes. Specifically, "
+                f"{non_geometries=} and {attributes=}"
             )
-        for id, att_name in zip(identifiers, attributes):
+        for o, att_name in zip(non_geometries, attributes):
             # attributes are a 1-indexed list, and Geometry is the first element, so we start from 2
             # while f'Output_{i}_attribute_name' not in
-            mod[id + "_attribute_name"] = att_name
-        os = [
-            outputs[i]
-            for i in range(len(outputs))
-            if outputs[i].socket_type != "NodeSocketGeometry"
-        ]
-        for o, domain in zip(os, domains):
+            mod[o.identifier + "_attribute_name"] = att_name
+        for o, domain in zip(non_geometries, domains):
             o.attribute_domain = domain
 
-        inputs = [s for s in mod.node_group.interface.items_tree if s.in_out == "INPUT"]
+        inputs = ng_inputs(mod.node_group)
         if not any(att_name is None for att_name in input_attributes):
             raise Exception("None should be provided for Geometry inputs.")
-        for i, att_name in zip(inputs, input_attributes):
-            id = i.identifier
+        for i, att_name in zip(inputs.values(), input_attributes):
+            o = i.identifier
             if att_name is not None:
-                mod[f"{id}_use_attribute"] = True
-                mod[f"{id}_attribute_name"] = att_name
+                mod[f"{o}_use_attribute"] = True
+                mod[f"{o}_attribute_name"] = att_name
 
     if apply:
         for obj in objs:

infinigen/core/util/blender.py

@@ -22,6 +22,7 @@
 
 from infinigen.core.nodes.node_info import DATATYPE_DIMS, DATATYPE_FIELDS
 
+from ..nodes.node_wrangler import ng_inputs
 from . import math as mutil
 from .logging import Suppress
 
@@ -555,11 +556,10 @@ def get_camera_res():
 def set_geomod_inputs(mod, inputs: dict):
     assert mod.type == "NODES"
     for k, v in inputs.items():
-        inputs = [s for s in mod.node_group.interface.items_tree if s.in_out == "INPUT"]
-        if not any(k == i.name for i in inputs):
+        inputs = ng_inputs(mod.node_group)
+        if k not in inputs:
             raise KeyError(f"Couldnt find {k=} in {mod.node_group.inputs.keys()=}")
-
-        soc = next(i for i in inputs if k == i.name)
+        soc = inputs[k]
 
         if not hasattr(soc, "default_value"):
             if v is not None:

infinigen/terrain/surface_kernel/kernelizer.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 
+from infinigen.core.nodes.node_wrangler import ng_inputs, ng_outputs
 from infinigen.terrain.utils import (
     NODE_ATTRS_AVAILABLE,
     NODE_FUNCTIONS,
@@ -48,7 +49,7 @@ def my_getattr(x, a):
 class Kernelizer:
     def get_inputs(self, node_tree):
         inputs = OrderedDict()
-        for node_input in node_tree.inputs:
+        for node_input in ng_inputs(node_tree).values():
             if node_input.type != SocketType.Geometry:
                 assert node_input.type != SocketType.Image
                 inputs[node_input.identifier] = SOCKETTYPE_KERNEL[node_input.type]
@@ -59,7 +60,7 @@ def get_output(self, node_tree):
         for node in node_tree.nodes:
             if node.bl_idname == Nodes.SetPosition:
                 outputs[Vars.Offset] = KernelDataType.float3
-        for node_output in node_tree.outputs:
+        for node_output in ng_outputs(node_tree).values():
             if node_output.type != SocketType.Geometry:
                 outputs[node_output.identifier] = SOCKETTYPE_KERNEL[node_output.type]
         return outputs
@@ -416,7 +417,7 @@ def __call__(self, modifier):
         code, imp_inputs, outputs = self.execute_node_tree(
             node_tree, collective_style=True
         )
-        for nodeoutput in node_tree.outputs:
+        for nodeoutput in ng_outputs(node_tree).values():
             id = nodeoutput.identifier
             if id != "Output_1":  # not Geometry
                 code = re.sub(rf"\b{id}\b", modifier[f"{id}_attribute_name"], code)

infinigen/terrain/utils/mesh.py

@@ -16,6 +16,7 @@
 from infinigen.core.util.logging import Timer
 from infinigen.core.util.organization import Attributes
 
+from ...core.nodes.node_wrangler import ng_inputs
 from .camera import getK
 from .ctype_util import ASDOUBLE, ASINT, load_cdll
 from .kernelizer_util import ATTRTYPE_DIMS, ATTRTYPE_FIELDS, NPTYPEDIM_ATTR, Vars
@@ -427,7 +428,7 @@ def move_modifier(target_obj, m):
     with Timer(f"copying {m.name}"):
         modifier = target_obj.modifiers.new(m.name, "NODES")
         modifier.node_group = m.node_group
-        for i, inp in enumerate(modifier.node_group.inputs):
+        for i, inp in enumerate(ng_inputs(modifier.node_group).values()):
             if i > 0:
                 id = inp.identifier
                 modifier[f"{id}_attribute_name"] = inp.name