lightspeed.trex.ai_tools.widget#

  • SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

  • SPDX-License-Identifier: Apache-2.0

  • Licensed under the Apache License, Version 2.0 (the “License”);

  • you may not use this file except in compliance with the License.

  • You may obtain a copy of the License at

  • https://www.apache.org/licenses/LICENSE-2.0

  • Unless required by applicable law or agreed to in writing, software

  • distributed under the License is distributed on an “AS IS” BASIS,

  • WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  • See the License for the specific language governing permissions and

  • limitations under the License.

class lightspeed.trex.ai_tools.widget.AIToolsWindowExtension#

AI Tools Window Extension.

Provides a workspace window for integrating with ComfyUI workflows, allowing users to submit AI-powered jobs and apply results to the USD stage.

on_shutdown() → None#

Clean up resources and close the workspace window.

on_startup(ext_id: str) → None#

Initialize the extension and create the workspace window.

class lightspeed.trex.ai_tools.widget.ArtifactHandler(context_name: str = '')#

Base class for handling artifacts produced by ComfyUI workflows.

Subclasses implement the logic for ingesting and applying specific artifact types (textures, meshes, etc.) to the USD stage.

Handlers “claim” artifacts they can handle from a list. This allows flexible matching based on metadata, file extension, or any other criteria.

Note: These handlers are internal to ComfyJobApplyHandler and not exposed through a separate registry. To add new artifact types, add them to the ARTIFACT_HANDLERS list in ComfyJobApplyHandler.

abstract async apply(
prim_paths: list[str],
artifacts: list[OutputArtifact],
) → None#

Apply artifacts to the given prims.

abstract classmethod claim(
artifacts: list[OutputArtifact],
) → list[OutputArtifact]#

Claim artifacts this handler can process.

class lightspeed.trex.ai_tools.widget.ComfyEventType#

Event type marker for ComfyUI state changes used by Stage Manager listeners.

class lightspeed.trex.ai_tools.widget.ComfyJobApplyHandler(context_name: str = '')#

Apply handler for ComfyUI jobs.

This handler processes ComfyJob results by delegating to specialized ArtifactHandlers for different output types.

Artifact handlers are defined in ARTIFACT_HANDLERS and are checked in order. Each handler “claims” artifacts it can process, and unclaimed artifacts are logged as warnings.

async apply(
interface: QueueInterface,
job: Job,
) → None#

Apply the job results.

Parameters:

  • interface – The job queue interface for accessing job events/results.

  • job – The completed job to apply results from.

classmethod can_handle(
job: Job,
) → bool#

Check if this handler can process the given job type.

get_apply_context(
job: Job,
) → dict[str, Any]#

Get context information for tracking applied state.

Override this to provide context-specific information that determines whether the job has already been applied (e.g., current USD edit target).

Returns an empty dict by default, meaning “applied once = applied everywhere”.

class lightspeed.trex.ai_tools.widget.ComfyJobGenerator(producer: ~collections.abc.Callable[[], ~collections.abc.Iterator[~pxr.Usd.Prim]], comfy_url: str = <factory>, context_name: str = '')#

Generates and submits ComfyUI jobs based on selected prims.

Prims are grouped by unique input data paths (resolved from LazyValues) to avoid duplicate work - if multiple prims share the same input files, only one job is created and the results are applied to all prims in the group.

Jobs are named after the input file (not the prim path). When results are applied, ComfyJobApplyHandler processes the job’s OutputArtifacts and delegates to specialized handlers (TextureArtifactHandler, etc.).

submit(
interface: QueueInterface,
workflow: Workflow,
use_inputs_for_output_filename_prefix: bool = False,
) → None#

Submit jobs for all prims returned by the producer, grouped by unique input data paths.

class lightspeed.trex.ai_tools.widget.ConnectionState(value)#

Connection state for the ComfyUI server.

property button_text: str#

Text for the connect/disconnect button.

property status_text: str#

Text describing the current connection status.

class lightspeed.trex.ai_tools.widget.LazyValue(func: ~collections.abc.Callable[[~pxr.Usd.Prim], ~lightspeed.trex.ai_tools.widget.lazy_values.T], return_type: type[~lightspeed.trex.ai_tools.widget.lazy_values.T], label: str, fields: list[~lightspeed.trex.ai_tools.widget.comfy.Field] = <factory>, description: str = '')#

A lazy value that is evaluated at job submission time.

class lightspeed.trex.ai_tools.widget.OutputArtifact(
path: ~pathlib.Path,
node_id: str = '',
metadata: dict[str,
~typing.Any] = <factory>,
)#

Represents an output artifact from a ComfyUI workflow.

Contains the downloaded file path along with any metadata from the workflow that can be used by handlers to determine how to process the artifact.

get(
key: str,
default: Any | None = None,
) → Any#

Get a metadata value by key.

class lightspeed.trex.ai_tools.widget.TextureArtifactHandler(context_name: str = '')#

Handles texture artifacts from ComfyUI workflows.

async apply(
prim_paths: list[str],
artifacts: list[OutputArtifact],
) → None#

Apply artifacts to the given prims.

classmethod claim(
artifacts: list[OutputArtifact],
) → list[OutputArtifact]#

Claim artifacts this handler can process.

class lightspeed.trex.ai_tools.widget.TexturePath(
texture_type: TextureTypes = TextureTypes.DIFFUSE,
)#

Lazy value that extracts a texture path from a prim.

Traverses related prims to find the shader and extract the texture path for the specified texture type.

texture_type#

The type of texture to extract (default: DIFFUSE).

Type:

omni.flux.asset_importer.core.data_models.enums.TextureTypes

Raises:

ValueError – If the prim doesn’t have exactly one texture of the specified type.

lightspeed.trex.ai_tools.widget.get_comfy_interface() → ComfyInterface#

Get the singleton ComfyInterface instance.

lightspeed.trex.ai_tools.widget.get_comfy_url() → str#

Get the current ComfyUI URL from settings.

Returns the persisted value if set, otherwise returns the default URL.

Iterate over prims related to the given prim.

Traverses the hierarchy to find related prims, following RTX Remix conventions: - Children of the prim - Prototypes (resolved from instances via get_prototype) - Materials (from mesh bindings) - Shaders (from materials)

This is used to find the shader prim for applying texture changes when the user selects an instance, mesh, or material.

Parameters:

prim – The starting prim.

Yields:

The prim itself and all related prims.

lightspeed.trex.ai_tools.widget.iter_selected_prims(
context_name: str = '',
) → Iterator[Prim]#

Iterate over currently selected prims in the USD stage.

Parameters:

context_name – The USD context name. Empty string uses the default context.

Yields:

Selected prims that exist in the stage.

Raises:

RuntimeError – If no USD stage is loaded.

lightspeed.trex.ai_tools.widget.iter_texture_path(
prim: Prim,
texture_type: TextureTypes = TextureTypes.DIFFUSE,
) → Iterator[Path]#

Iterate over texture paths for a given prim and texture type.

Traverses related prims to find shader prototypes, then extracts texture paths for the specified texture type.

Parameters:

  • prim – The prim to find textures for.

  • texture_type – The type of texture to find (default: DIFFUSE).

Yields:

Paths to texture files.

lightspeed.trex.ai_tools.widget.set_comfy_url(url: str) → None#

Set the ComfyUI URL in persistent settings.

This value will be remembered between sessions.

lightspeed.trex.ai_tools.widget.submit_selected_prims_to_comfy(
interface: QueueInterface | None = None,
workflow: Workflow | None = None,
use_inputs_for_output_filename_prefix: bool = False,
) → None#

Submit the currently selected prims to ComfyUI for processing.

This is a convenience function that creates a ComfyJobGenerator with the current selection and submits jobs using the provided (or default) interface and workflow.

Parameters:

  • interface – The job queue interface. If None, uses the default interface.

  • workflow – The workflow to use. If None, uses the workflow from the ComfyInterface.

  • use_inputs_for_output_filename_prefix – If True, uses input filenames for output prefixes.