From 419f2f4ef7cd9bf62364151b222562065467ac58 Mon Sep 17 00:00:00 2001 From: Lingjie Mei Date: Mon, 16 Sep 2024 03:01:34 -0400 Subject: [PATCH] Fix inputs --- .../objects/creatures/parts/ridged_fin.py | 14 ++----- .../objects/creatures/util/part_util.py | 11 ++--- infinigen/assets/objects/seating/sofa.py | 2 +- .../assets/utils/extract_nodegroup_parts.py | 8 ++-- .../core/nodes/node_transpiler/transpiler.py | 3 +- infinigen/core/nodes/node_wrangler.py | 13 ++++-- infinigen/core/surface.py | 40 ++++++++----------- infinigen/core/util/blender.py | 8 ++-- .../terrain/surface_kernel/kernelizer.py | 7 ++-- infinigen/terrain/utils/mesh.py | 3 +- 10 files changed, 52 insertions(+), 57 deletions(-) diff --git a/infinigen/assets/objects/creatures/parts/ridged_fin.py b/infinigen/assets/objects/creatures/parts/ridged_fin.py index 60f7dd7f1..20d001ca5 100644 --- a/infinigen/assets/objects/creatures/parts/ridged_fin.py +++ b/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) diff --git a/infinigen/assets/objects/creatures/util/part_util.py b/infinigen/assets/objects/creatures/util/part_util.py index 49b5b5167..b61c205fb 100644 --- a/infinigen/assets/objects/creatures/util/part_util.py +++ b/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 diff --git a/infinigen/assets/objects/seating/sofa.py b/infinigen/assets/objects/seating/sofa.py index 1274b52df..ee0516dce 100644 --- a/infinigen/assets/objects/seating/sofa.py +++ b/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) diff --git a/infinigen/assets/utils/extract_nodegroup_parts.py b/infinigen/assets/utils/extract_nodegroup_parts.py index afa6e5e0c..8062e967b 100644 --- a/infinigen/assets/utils/extract_nodegroup_parts.py +++ b/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) diff --git a/infinigen/core/nodes/node_transpiler/transpiler.py b/infinigen/core/nodes/node_transpiler/transpiler.py index 80b23f9e7..79a58b484 100644 --- a/infinigen/core/nodes/node_transpiler/transpiler.py +++ b/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") diff --git a/infinigen/core/nodes/node_wrangler.py b/infinigen/core/nodes/node_wrangler.py index 823b9c3ba..602592dbf 100644 --- a/infinigen/core/nodes/node_wrangler.py +++ b/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] diff --git a/infinigen/core/surface.py b/infinigen/core/surface.py index b7e3d56b0..492d29131 100644 --- a/infinigen/core/surface.py +++ b/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: diff --git a/infinigen/core/util/blender.py b/infinigen/core/util/blender.py index 5a60d048f..c22f88b51 100644 --- a/infinigen/core/util/blender.py +++ b/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: diff --git a/infinigen/terrain/surface_kernel/kernelizer.py b/infinigen/terrain/surface_kernel/kernelizer.py index 95ac6cbd8..2db9dc7d4 100644 --- a/infinigen/terrain/surface_kernel/kernelizer.py +++ b/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) diff --git a/infinigen/terrain/utils/mesh.py b/infinigen/terrain/utils/mesh.py index 77b2dd9f6..bc390f216 100644 --- a/infinigen/terrain/utils/mesh.py +++ b/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