procfunc 0.30.0__py3-none-any.whl

procfunc/control.py

@@ -0,0 +1,174 @@
+import functools
+import logging
+from typing import Callable, TypeVar
+
+import numpy as np
+
+import procfunc.compute_graph as cg
+from procfunc import context
+from procfunc import types as t
+from procfunc.tracer import (
+    PatchFunctionTarget,
+    RngProxy,
+    TraceLevel,
+    register_trace_target,
+)
+from procfunc.util import pytree
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def _unwrap_proxy(x):
+    return x.node if isinstance(x, cg.Proxy) else x
+
+
+def choice_idx(rng: np.random.Generator, weights: list[float]) -> int:
+    weights = np.array(weights)
+    weights = weights / weights.sum()
+    return rng.choice(len(weights), p=weights)
+
+
+def _peekthrough_execute_all_choices(
+    choice_options: list[tuple[Callable[..., T], float]],
+    chosen_idx: int,
+    args: tuple,
+    kwargs: dict,
+) -> cg.Proxy:
+    evald = []
+    for func, weight in choice_options:
+        res = func(*args, **kwargs)
+        res = pytree.PyTree(res).map(_unwrap_proxy).obj()
+        evald.append((res, weight))
+
+    new_kwargs = dict(
+        choice_options=evald,
+        chosen_idx=chosen_idx,
+    )
+    node = cg.FunctionCallNode(
+        func=choice,
+        args=(),
+        kwargs=new_kwargs,
+    )
+    return cg.Proxy(node)
+
+
+class ChoiceResultProxy(cg.Proxy):
+    def __init__(self, node: cg.FunctionCallNode):
+        assert node.func is choice
+        super().__init__(node)
+
+    def __call__(self, *args, **kwargs) -> cg.Proxy:
+        current_level = context.globals.current_trace_level
+        if current_level is None:
+            raise ValueError(f"Executed {self} while not in a tracing context?")
+        if current_level >= TraceLevel.RANDOM_CONTROL:
+            return _peekthrough_execute_all_choices(
+                self.node.kwargs["choice_options"],
+                self.node.kwargs["chosen_idx"],
+                args,
+                kwargs,
+            )
+        else:
+            idx = self.node.kwargs["chosen_idx"]
+            chosen_value = self.node.kwargs["choice_options"][idx][0]
+            return chosen_value(*args, **kwargs)
+
+
+def _choice_create_custom_tracer_wrapper(
+    target: PatchFunctionTarget,
+    patcher,
+    **_kwargs,
+):
+    """
+    choice needs a custom tracing wrapper since we want to trace what happens inside its sub-options.
+
+    Behavior depends on trace_level vs RANDOM_CONTROL:
+        - >= RANDOM_CONTROL: peekthrough all options exhaustively
+        - < RANDOM_CONTROL: trace only the branch that would actually be chosen (active resolution)
+    """
+
+    @functools.wraps(choice)
+    def wrapper(
+        choice_rng: RngProxy,
+        choice_options: list[tuple[T, float]] | None = None,
+        chosen_idx: int | None = None,
+    ):
+        choice_values, choice_weights = zip(*choice_options)
+        if chosen_idx is None:
+            chosen_idx = choice_idx(choice_rng.rng, choice_weights)
+
+        unwrapped_options = [
+            (pytree.PyTree(v).map(_unwrap_proxy).obj(), w) for v, w in choice_options
+        ]
+        new_kwargs = dict(
+            choice_rng=choice_rng.node,
+            choice_options=unwrapped_options,
+            chosen_idx=chosen_idx,
+        )
+        node = cg.FunctionCallNode(
+            func=choice,
+            args=(),
+            kwargs=new_kwargs,
+        )
+        return ChoiceResultProxy(node)
+
+    return wrapper
+
+
+TArgs = TypeVar("TArgs")
+
+
+def choice(
+    choice_rng: np.random.Generator,
+    choice_options: list[tuple[T, float]],
+    chosen_idx: int | None = None,
+) -> T:
+    """
+    Args:
+        rng: random number generator
+        weights: list of weights for each option - will be normalized to sum to 1 as a probability distribution
+        options: list of callables to choose from
+        chosen: if not None, use this value instead of choosing randomly. Cannot be traced, but may appear in the output of the tracer.
+        chosen_idx: if not None, use this index instead of choosing randomly. Cannot be traced, but may appear in the output of the tracer.
+        **child_kwargs: keyword arguments to pass to the chosen callable
+
+    TODO: this function should really take args and kwargs as tuple & dict, not via expansion
+    but this requires our Node to handle cases with pytrees as inputs. Currently it would fail to execute the children if they are hidden in a pytree
+    """
+
+    choice_values, choice_weights = zip(*choice_options)
+    if chosen_idx is None:
+        chosen_idx = choice_idx(choice_rng, choice_weights)
+    option = choice_values[chosen_idx]
+    return option
+
+
+_ORIG_EXECUTE_CHOICE = choice
+register_trace_target(
+    func=choice,
+    trace_level=TraceLevel.RANDOM_CONTROL,
+    allow_exec=False,
+    custom_trace_wrapper_create=_choice_create_custom_tracer_wrapper,
+)
+
+
+def sample_collection(
+    rng: np.random.Generator,
+    func: Callable[[np.random.Generator], T],
+    n: int,
+    skip_none: bool = False,
+) -> t.Collection:
+    rngs = rng.spawn(n)
+    objects = [func(rng) for rng in rngs]
+    if skip_none:
+        objects = [obj for obj in objects if obj is not None]
+    return t.Collection(objects, name=func.__name__)
+
+
+__all__ = [
+    "choice",
+    "choice_idx",
+    "choice",
+]

procfunc/nodes/__init__.py

@@ -0,0 +1,66 @@
+from pandas import read_json as _read_json
+
+from procfunc.tracer import autowrap_module as _autowrap
+from procfunc.util.manifest import module_path
+
+from . import compositor, func, geo, math, shader
+from .bpy_node_info import NodeDataType, NodeGroupType, SocketType
+
+# ruff: noqa: E402
+_autowrap(compositor, allow_exec=False)
+_autowrap(func, allow_exec=False)
+_autowrap(geo, allow_exec=False)
+_autowrap(math, allow_exec=False)
+_autowrap(shader, allow_exec=False)
+
+from .execute.execute import (
+    as_nodegroup,
+    to_aliases,
+    to_compositor,
+    to_curve_object,
+    to_environment,
+    to_light,
+    to_mesh_object,
+    to_mesh_object_with_attributes,
+    to_objects_multi,
+)
+from .execute.util import NODE_OPERATOR_TABLE
+from .node_function import node_function
+from .types import (
+    ProcNode,
+    Shader,
+    SocketOrVal,
+)
+
+NODES_MANIFEST_PATH = module_path() / "nodes" / "manifest.json"
+assert NODES_MANIFEST_PATH.exists(), f"Manifest not found at {NODES_MANIFEST_PATH}"
+NODES_MANIFEST = _read_json(NODES_MANIFEST_PATH)
+
+__all__ = [
+    # Node category submodules
+    "compositor",
+    "func",
+    "geo",
+    "math",
+    "shader",
+    # Core types
+    "ProcNode",
+    "Shader",
+    "SocketOrVal",
+    "NodeDataType",
+    "NodeGroupType",
+    "SocketType",
+    # Graph execution
+    "as_nodegroup",
+    "to_aliases",
+    "to_compositor",
+    "to_curve_object",
+    "to_environment",
+    "to_light",
+    "to_material",
+    "to_mesh_object",
+    "to_mesh_object_with_attributes",
+    "to_objects_multi",
+    # User-facing decorator for custom node functions
+    "node_function",
+]

procfunc/nodes/bindings_util.py

@@ -0,0 +1,196 @@
+import logging
+from dataclasses import dataclass
+from enum import Enum
+from typing import Self
+
+import procfunc as pf
+from procfunc.nodes.bpy_node_info import NodeDataType, NodeGroupType
+from procfunc.util.log import raise_error_or_warn
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RuntimeResolveDataType:
+    data_types: list[NodeDataType]
+    dependent_input_names: list[str]
+
+
+class ContextualNode(Enum):
+    """Special node types that map to different implementations based on context."""
+
+    COMBINE_COLOR = "ContextualCombineColor"
+    SEPARATE_COLOR = "ContextualSeparateColor"
+    GROUP = "ContextualGroup"
+    MIX_RGB = "ContextualMixRGB"
+    MATH = "ContextualMath"
+
+    @classmethod
+    def parse_name(cls, from_name: str) -> Self | None:
+        try:
+            return cls(from_name)
+        except ValueError:
+            return None
+
+
+@dataclass
+class NodeContextResolution:
+    contextual_node: ContextualNode
+    node_group_type: NodeGroupType
+    node_type: str
+    input_keys_map: dict[str, tuple[str, int] | str | None] | None
+
+
+# Mapping to map unified version of nodes to context-specific blender versions
+# Input keys are (ContextualNode, NodeGroupType) - the unifiied version + the context we need to use it in
+# Output is the (NodeType, input map) - tells what specific node to construct, then how to change the inputs for that context, if applicable
+
+CONTEXTUAL_NODE_MAPPING = [
+    NodeContextResolution(
+        contextual_node=ContextualNode.COMBINE_COLOR,
+        node_group_type=NodeGroupType.SHADER,
+        node_type="ShaderNodeCombineColor",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.COMBINE_COLOR,
+        node_group_type=NodeGroupType.TEXTURE,
+        node_type="TextureNodeCombineColor",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.COMBINE_COLOR,
+        node_group_type=NodeGroupType.GEOMETRY,
+        node_type="FunctionNodeCombineColor",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.SEPARATE_COLOR,
+        node_group_type=NodeGroupType.SHADER,
+        node_type="ShaderNodeSeparateColor",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.SEPARATE_COLOR,
+        node_group_type=NodeGroupType.TEXTURE,
+        node_type="TextureNodeSeparateColor",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.SEPARATE_COLOR,
+        node_group_type=NodeGroupType.GEOMETRY,
+        node_type="FunctionNodeSeparateColor",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.SEPARATE_COLOR,
+        node_group_type=NodeGroupType.COMPOSITOR,
+        node_type="CompositorNodeSeparateColor",
+        input_keys_map={"Color": "Image"},
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.GROUP,
+        node_group_type=NodeGroupType.COMPOSITOR,
+        node_type="CompositorNodeGroup",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.GROUP,
+        node_group_type=NodeGroupType.TEXTURE,
+        node_type="TextureNodeGroup",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.GROUP,
+        node_group_type=NodeGroupType.GEOMETRY,
+        node_type="GeometryNodeGroup",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.MIX_RGB,
+        node_group_type=NodeGroupType.SHADER,
+        node_type="ShaderNodeMixRGB",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.MIX_RGB,
+        node_group_type=NodeGroupType.TEXTURE,
+        node_type="TextureNodeMixRGB",
+        input_keys_map={"Fac": "Factor"},
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.MIX_RGB,
+        node_group_type=NodeGroupType.GEOMETRY,
+        node_type="ShaderNodeMixRGB",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.MATH,
+        node_group_type=NodeGroupType.SHADER,
+        node_type="ShaderNodeMath",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.MATH,
+        node_group_type=NodeGroupType.COMPOSITOR,
+        node_type="CompositorNodeMath",
+        input_keys_map=None,
+    ),
+    NodeContextResolution(
+        contextual_node=ContextualNode.MATH,
+        node_group_type=NodeGroupType.TEXTURE,
+        node_type="TextureNodeMath",
+        input_keys_map=None,
+    ),
+]
+
+
+def resolve_contextual_node(
+    node: list[NodeContextResolution], group: NodeGroupType
+) -> tuple[str, dict[str, str]]:
+    item = next(
+        (
+            item
+            for item in CONTEXTUAL_NODE_MAPPING
+            if (item.node_group_type == group and item.contextual_node == node)
+        ),
+        None,
+    )
+
+    if item is None:
+        raise ValueError(f"No contextual node found for {node} in {group}")
+
+    return item.node_type, item.input_keys_map
+
+
+def raise_shader_normal_error(node_func_name: str, logger: logging.Logger):
+    """Helper to handle normal socket usage in shaders based on context setting."""
+
+    message = (
+        f"Using 'normal' input in {node_func_name} is not recommended. "
+        f"Use displacement instead. To suppress this message, set "
+        f"pf.context.globals.warn_mode_avoid_normal_bump = 'ignore'"
+    )
+
+    raise_error_or_warn(message, pf.context.globals.warn_mode_avoid_normal_bump, logger)
+
+
+def raise_explicit_noise_vector_error(node_func_name: str, logger: logging.Logger):
+    """Helper to handle missing vector input in noise functions based on context setting."""
+    message = (
+        f"The 'vector' input is required for {node_func_name}. "
+        f"Infinigen requires an explicit vector input - node will not default to using texture coordinate or world coordinate like blender. "
+        f"To suppress this message, set pf.context.globals.warn_mode_avoid_implicit_vector = 'ignore'"
+    )
+    raise_error_or_warn(
+        message, pf.context.globals.warn_mode_avoid_implicit_vector, logger
+    )
+
+
+def raise_io_error(node_func_name: str, logger: logging.Logger):
+    """Helper to handle IO nodes based on context setting."""
+    message = (
+        f"IO nodes are not allowed in nodegroups but we found {node_func_name} in the nodegroup. Use the nodegroup interface instead. "
+        "To suppress this message, set pf.context.globals.warn_mode_avoid_io_nodes = 'ignore'"
+    )
+    raise_error_or_warn(message, pf.context.globals.warn_mode_avoid_io_nodes, logger)