nodes/NVIDIAGameWorks-ComfyUI-RTX-Remix
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implemented Remix Save Node + Move API Nodes

Browse Source
This commit is contained in:
Pierre-Olivier Trottier
2025-11-12 16:54:37 -08:00
parent a9e8f4fedc
commit c771dfb6e9
9 changed files with 304 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/__init__.py
+17 -14
View File
@@ -15,19 +15,9 @@
* limitations under the License.
"""
from .common import (
EndContext,
InvertBool,
RestAPIDetails,
StartContext,
StringConcatenate,
StringConstant,
StrToList,
Switch,
)
from .file import DeleteFile
from .ingestion import GetDefaultDirectory, IngestTexture
from .layers import (
from .api.file import DeleteFile
from .api.ingestion import GetDefaultDirectory, IngestTexture
from .api.layers import (
CloseProject,
CreateLayer,
DefineLayerId,
@@ -42,13 +32,24 @@ from .layers import (
SaveLayer,
SetEditTarget,
)
from .textures import (
from .api.textures import (
GetTextures,
SetTexture,
TexturesType,
TexturesTypes,
TextureTypeToUSDAttribute,
)
from .common import (
EndContext,
InvertBool,
RestAPIDetails,
StartContext,
StringConcatenate,
StringConstant,
StrToList,
Switch,
)
from .save import RemixSaveTexture
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
@@ -71,6 +72,7 @@ NODE_CLASS_MAPPINGS = {
"RTXRemixOpenProject": OpenProject,
"RTXRemixRemoveLayer": RemoveLayer,
"RTXRemixRestAPIDetails": RestAPIDetails,
"RTXRemixSaveTexture": RemixSaveTexture,
"RTXRemixSaveLayer": SaveLayer,
"RTXRemixSetEditTarget": SetEditTarget,
"RTXRemixSetTexture": SetTexture,
@@ -104,6 +106,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"RTXRemixOpenProject": "RTX Remix Open Project",
"RTXRemixRemoveLayer": "RTX Remix Remove Layer",
"RTXRemixRestAPIDetails": "RTX Remix Rest API Details",
"RTXRemixSaveTexture": "RTX Remix Save Texture",
"RTXRemixSaveLayer": "RTX Remix Save Layer",
"RTXRemixSetEditTarget": "RTX Remix Set Edit Target",
"RTXRemixSetTexture": "RTX Remix Set Texture",
nodes/file.py → nodes/api/file.py
+3 -3
View File
@@ -17,8 +17,8 @@
import pathlib
from .common import add_context_input_enabled_and_output
from .constant import PREFIX_MENU
from ..common import add_context_input_enabled_and_output
from ..constant import PREFIX_MENU_API
_file_name = pathlib.Path(__file__).stem
@@ -41,7 +41,7 @@ class DeleteFile:
RETURN_TYPES = ("BOOL",)
RETURN_NAMES = ("File deleted",)
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def execute(self, path: str):
if not self.enable_this_node: # noqa
nodes/ingestion.py → nodes/api/ingestion.py
+5 -5
View File
@@ -25,9 +25,9 @@ import requests
import torch
from PIL import Image
from .common import add_context_input_enabled_and_output
from .constant import HEADER_LSS_REMIX_VERSION_1_0, PREFIX_MENU
from .utils import check_response_status_code, posix
from ..common import add_context_input_enabled_and_output
from ..constant import HEADER_LSS_REMIX_VERSION_1_0, PREFIX_MENU_API
from ..utils import check_response_status_code
_file_name = pathlib.Path(__file__).stem
@@ -75,7 +75,7 @@ class IngestTexture:
RETURN_NAMES = ("texture_path",)
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def ingest_texture(
self,
@@ -162,7 +162,7 @@ class GetDefaultDirectory:
OUTPUT_NODE = False
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def _get_default_output_directory(self) -> str:
"""Utility method to get default output directory from RTX Remix API."""
nodes/layers.py → nodes/api/layers.py
+14 -19
View File
@@ -24,13 +24,13 @@ from urllib.parse import quote, unquote
import requests
from .common import (
from ..common import (
RemixContext,
add_context_input_enabled_and_output,
get_context_inputs,
)
from .constant import HEADER_LSS_REMIX_VERSION_1_0, PREFIX_MENU
from .utils import check_response_status_code, merge_dict, posix
from ..constant import HEADER_LSS_REMIX_VERSION_1_0, PREFIX_MENU_API
from ..utils import check_response_status_code, merge_dict, posix
NONE = "None"
_layer_types = [
@@ -100,7 +100,7 @@ class DefineLayerId:
RETURN_NAMES = ("layer_id",)
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def execute(
self,
@@ -111,7 +111,7 @@ class DefineLayerId:
# Handle None case for parent_layer_id
if parent_layer_id is None:
parent_layer_id = ""
layer_id_dir = pathlib.Path(parent_layer_id).parent
if directories:
layer_id_dir = layer_id_dir / directories
@@ -162,7 +162,7 @@ class CreateLayer:
OUTPUT_NODE = False
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def create_layer(
self,
@@ -199,7 +199,6 @@ class LayerType:
@classmethod
def INPUT_TYPES(cls): # noqa N802
inputs = {
"required": {
"layer_type": (_layer_types,),
@@ -212,7 +211,7 @@ class LayerType:
FUNCTION = "get_layer_type"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_layer_type(self, layer_type: str) -> tuple[str]:
if not self.enable_this_node: # noqa
@@ -227,7 +226,6 @@ class LayerTypes:
@classmethod
def INPUT_TYPES(cls): # noqa N802
inputs = {
"required": {
"layer_types": (
@@ -246,7 +244,7 @@ class LayerTypes:
FUNCTION = "get_layer_types"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_layer_types(self, layer_types: str) -> tuple[str]:
if not self.enable_this_node: # noqa
@@ -306,7 +304,7 @@ class GetLayers:
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def execute(
self,
@@ -320,7 +318,7 @@ class GetLayers:
if not self.enable_this_node: # noqa
# Return non-empty lists even when disabled to avoid ComfyUI empty list bug
return ([""], [""], False)
layer_types_list = [t.strip() for t in layer_types.split(",")]
params = {
"layer_types": layer_types_list,
@@ -381,10 +379,8 @@ class GetLayers:
class _LayerOp:
@classmethod
def INPUT_TYPES(cls): # noqa N802
inputs = {
"required": {
"layer_id": (
@@ -400,7 +396,7 @@ class _LayerOp:
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
@abc.abstractmethod
def execute(self, layer_id: str) -> tuple[str]:
@@ -447,7 +443,6 @@ class RemoveLayer(_LayerOp):
@classmethod
def INPUT_TYPES(cls): # noqa N802
inputs = {
"required": {
"parent_layer_id": (
@@ -503,7 +498,7 @@ class GetEditTarget:
OUTPUT_NODE = False
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_edit_target(self, context: RemixContext) -> tuple[str]:
if not self.enable_this_node: # noqa
@@ -557,7 +552,7 @@ class CloseProject:
OUTPUT_NODE = False
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def close_project(self, force: bool = False) -> tuple[str]:
if not self.enable_this_node: # noqa
@@ -605,7 +600,7 @@ class GetLoadedProject:
OUTPUT_NODE = False
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_loaded_project(self) -> tuple[str]:
if not self.enable_this_node: # noqa
nodes/textures.py → nodes/api/textures.py
+13 -18
View File
@@ -23,9 +23,9 @@ import requests
import torch
from PIL import Image, ImageOps
from .common import add_context_input_enabled_and_output
from .constant import HEADER_LSS_REMIX_VERSION_1_0, PREFIX_MENU
from .utils import check_response_status_code, posix
from ..common import add_context_input_enabled_and_output
from ..constant import HEADER_LSS_REMIX_VERSION_1_0, PREFIX_MENU_API
from ..utils import check_response_status_code, posix
_texture_types = [
"DIFFUSE",
@@ -130,7 +130,7 @@ class GetTextures:
)
FUNCTION = "get_texture_prims_assets"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_texture_prims_assets(
self,
@@ -141,7 +141,6 @@ class GetTextures:
layer_id: str | None = None,
exists: bool = True,
) -> tuple[list[str], list[str], list[torch.Tensor], list[torch.Tensor]]:
if not self.enable_this_node: # noqa
return ([], [], [], [])
@@ -162,7 +161,7 @@ class GetTextures:
textures = json.loads(r.text).get("textures", [])
if not textures:
raise ValueError(
"No textures found. Please check the parameters of your node.\n" f"URL: {r.url}, PARAMS: {payload}"
f"No textures found. Please check the parameters of your node.\nURL: {r.url}, PARAMS: {payload}"
)
result_attrs = []
@@ -180,10 +179,10 @@ class GetTextures:
result_images.append(image)
texture_names.append(pathlib.Path(texture_path).stem)
result_attrs.append(usd_attr)
if 'A' in img_1.getbands():
mask = np.array(img_1.getchannel('A')).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask)
if "A" in img_1.getbands():
mask = np.array(img_1.getchannel("A")).astype(np.float32) / 255.0
mask = 1.0 - torch.from_numpy(mask)
else:
mask = torch.zeros((image.shape[2], image.shape[3]), dtype=torch.float32, device="cpu")
result_masks.append(mask.unsqueeze(0))
@@ -207,7 +206,6 @@ class TexturesTypes:
@classmethod
def INPUT_TYPES(cls): # noqa N802
inputs = {
"required": {
"texture_types": (
@@ -226,7 +224,7 @@ class TexturesTypes:
RETURN_NAMES = ("texture_types",)
FUNCTION = "get_texture_types"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_texture_types(self, texture_types: str) -> tuple[str]:
if not self.enable_this_node: # noqa
@@ -243,7 +241,6 @@ class TexturesType:
@classmethod
def INPUT_TYPES(cls): # noqa N802
inputs = {
"required": {
"texture_type": (_texture_types,),
@@ -255,7 +252,7 @@ class TexturesType:
RETURN_NAMES = ("texture_type",)
FUNCTION = "get_texture_type"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_texture_type(self, texture_type: str) -> tuple[str]:
if not self.enable_this_node: # noqa
@@ -289,10 +286,9 @@ class SetTexture:
RETURN_TYPES = ()
RETURN_NAMES = ()
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def set_texture(self, usd_attribute: str, texture_path: str, force: bool = False):
if not self.enable_this_node: # noqa
return ()
@@ -336,10 +332,9 @@ class TextureTypeToUSDAttribute:
RETURN_NAMES = ("usd_attribute",)
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_attr_from_texture_type(self, usd_attribute: str, texture_type: str):
if not self.enable_this_node: # noqa
return ("",)
nodes/common.py
+9 -9
View File
@@ -18,7 +18,7 @@
import pathlib
from collections import namedtuple
from .constant import CONTEXT_TYPE, PREFIX_MENU
from .constant import CONTEXT_TYPE, PREFIX_MENU_API
from .utils import merge_dict
_file_name = pathlib.Path(__file__).stem
@@ -125,7 +125,7 @@ class RestAPIDetails:
FUNCTION = "get_address"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_address(self, address, port):
return address, port
@@ -144,7 +144,7 @@ class StartContext:
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def execute(self, address, port):
return (RemixContext(address, port),)
@@ -162,7 +162,7 @@ class EndContext:
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
OUTPUT_NODE = True
@@ -183,7 +183,7 @@ class StringConstant:
RETURN_TYPES = ("STRING",)
FUNCTION = "get_string"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def get_string(self, string):
return (string,)
@@ -206,7 +206,7 @@ class StringConcatenate:
RETURN_TYPES = ("STRING",)
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def execute(self, string1, string2, separator="_"):
return (string1 + separator + string2,)
@@ -241,7 +241,7 @@ class Switch:
INPUT_IS_LIST = True # need this or it will crash
OUTPUT_IS_LIST = (True,)
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def execute(self, if_true: _any, if_false: _any, switcher: list[bool]):
return (if_true if switcher[0] else if_false,)
@@ -263,7 +263,7 @@ class InvertBool:
RETURN_TYPES = ("BOOLEAN",)
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def execute(self, value: bool):
return (not value,)
@@ -286,7 +286,7 @@ class StrToList:
RETURN_TYPES = ("STRING",)
OUTPUT_IS_LIST = (True,)
FUNCTION = "execute"
CATEGORY = f"{PREFIX_MENU}/{_file_name}"
CATEGORY = f"{PREFIX_MENU_API}/{_file_name}"
def execute(self, value: str):
return ([value],)
nodes/constant.py
+4 -1
View File
@@ -15,6 +15,9 @@
* limitations under the License.
"""
PREFIX_MENU = "RTXRemix"
PREFIX_BASE = "RTX Remix"
PREFIX_MENU_API = f"{PREFIX_BASE}/API"
HEADER_LSS_REMIX_VERSION_1_0 = {"Accept": "application/lightspeed.remix.service+json; version=1.0"}
CONTEXT_TYPE = "RTXRemixContext"
nodes/save.py
+234
View File
@@ -0,0 +1,234 @@
"""
* 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.
"""
from __future__ import annotations
__all__ = ["TextureType", "FileExtension", "Colorspace", "RemixSaveTexture"]
import uuid
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING
import folder_paths
import numpy as np
from comfy_execution.utils import get_executing_context
from Imath import Channel, PixelType
from OpenEXR import Header, OutputFile
from PIL import Image
from .constant import PREFIX_BASE
if TYPE_CHECKING:
import torch
# Cache for current execution timestamp (ComfyUI processes prompts sequentially)
_prompt_id = None
_prompt_timestamp = None
class TextureType(Enum):
DIFFUSE = "albedo"
ROUGHNESS = "roughness"
ANISOTROPY = "anisotropy"
METALLIC = "metallic"
EMISSIVE = "emissive_mask"
NORMAL_OGL = "normal_ogl"
NORMAL_DX = "normal_dx"
NORMAL_OTH = "normal_oth"
HEIGHT = "height"
TRANSMITTANCE = "transmittance"
MEASUREMENT_DISTANCE = "measurement_distance"
SINGLE_SCATTERING = "single_scattering"
OTHER = "other"
class FileExtension(Enum):
PNG = ".png"
EXR = ".exr"
class Colorspace(Enum):
LINEAR = "linear"
SRGB = "sRGB"
class RemixSaveTexture:
"""Save a texture to the file system"""
@classmethod
def INPUT_TYPES(cls): # noqa N802
inputs = {
"required": {
"textures": ("IMAGE", {}),
"texture_type": (
[e.value for e in TextureType],
{"default": TextureType.DIFFUSE.value},
),
"file_extension": (
[e.value for e in FileExtension],
{"default": FileExtension.PNG.value},
),
"colorspace": (
[e.value for e in Colorspace],
{"default": Colorspace.LINEAR.value},
),
},
"hidden": {
"extra_pnginfo": "EXTRA_PNGINFO",
},
}
return inputs
DESCRIPTION = "Save a texture to the file system to be used in RTX Remix"
FUNCTION = "save_image"
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("file_path",)
CATEGORY = PREFIX_BASE
OUTPUT_NODE = True
def save_image(
self,
textures: torch.Tensor,
texture_type: str,
file_extension: str,
colorspace: str,
extra_pnginfo=None,
):
global _prompt_id, _prompt_timestamp
remix_job_id = None
# Check if metadata is stored in extra_pnginfo
if extra_pnginfo is not None:
remix_job_id = extra_pnginfo.get("rtx-remix", {}).get("job_id")
# Get the execution context to access the unique prompt_id for this run
context = get_executing_context()
prompt_id = context.prompt_id if context else None
# Create timestamp for new execution, reuse for same execution
timestamp = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
if prompt_id and prompt_id != _prompt_id:
# New execution - create new timestamp
_prompt_id = prompt_id
_prompt_timestamp = timestamp
# Use remix_job_id or current execution timestamp as subdirectory
sub_directory = remix_job_id or _prompt_timestamp
output_directory = Path(folder_paths.get_output_directory()) / sub_directory
# Create output directory if it doesn't exist
output_directory.mkdir(parents=True, exist_ok=True)
output_paths = []
results = []
for texture in textures:
output_folder, file_name, counter, _sub_folder, _file_name_prefix = folder_paths.get_save_image_path(
texture_type,
str(output_directory),
texture.shape[1],
texture.shape[0],
)
file_path = Path(output_folder) / f"{file_name}_{counter:05}_{file_extension}"
# ComfyUI works in linear space, convert to numpy
np_img = texture.cpu().numpy()
# Apply colorspace conversion if needed
# Note: ComfyUI internal format is linear [0,1]
if colorspace == Colorspace.SRGB.value:
# Convert from linear to sRGB
np_img = self._linear_to_srgb(np_img)
if file_extension == FileExtension.EXR.value:
# Save the image using OpenEXR
float_img = np_img.astype(np.float32)
height, width, _ = float_img.shape
# Prepare the EXR header
header = Header(width, height)
header["channels"] = dict([(c, Channel(PixelType(PixelType.FLOAT))) for c in "RGB"])
# Split the channels for OpenEXR
R = float_img[:, :, 0].tobytes()
G = float_img[:, :, 1].tobytes()
B = float_img[:, :, 2].tobytes()
# Write the EXR file
exr_file = OutputFile(str(file_path), header)
try:
exr_file.writePixels({"R": R, "G": G, "B": B})
finally:
exr_file.close()
# Create a preview PNG for the UI (browsers can't display EXR)
preview_filename = f"_temp_{uuid.uuid4().hex[:8]}_{file_name}_{counter:05}.png"
preview_path = Path(folder_paths.get_temp_directory()) / preview_filename
# Save preview as PNG
Image.fromarray(np.clip(255.0 * np_img, 0, 255).astype(np.uint8)).save(preview_path)
# Add preview for UI display
results.append(
{
"filename": preview_filename,
"subfolder": "",
"type": "temp",
}
)
else:
# Save the image using Pillow
Image.fromarray(np.clip(255.0 * np_img, 0, 255).astype(np.uint8)).save(file_path, compress_level=0)
# Add actual file for UI display
results.append(
{
"filename": file_path.name,
"subfolder": sub_directory,
"type": "output",
}
)
output_paths.append(str(file_path))
return {"ui": {"images": results}, "result": (output_paths,)}
def _linear_to_srgb(self, linear: np.ndarray) -> np.ndarray:
"""
Convert linear RGB to sRGB (apply gamma encoding).
ComfyUI works in linear RGB [0,1]. This applies the sRGB OETF to encode the image.
Formula: if x ≤ 0.0031308: y = 12.92x
else: y = 1.055x^(1/2.4) - 0.055
"""
return np.where(linear <= 0.0031308, linear * 12.92, 1.055 * np.power(linear, 1.0 / 2.4) - 0.055)
@classmethod
def IS_CHANGED(cls, **kwargs): # noqa N802
"""
Always process the node
"""
return float("nan")
requirements.txt
+5 -4
View File
@@ -1,4 +1,5 @@
numpy # SWIPAT filed under: https://nvbugs/4702469
pillow>=10.0.1 # SWIPAT filed under: https://nvbugs/4702475
requests # SWIPAT filed under: https://nvbugs/4702237
torch # SWIPAT filed under: https://nvbugs/4702252
numpy # OSRB filed under: https://nvbugs/4702469
pillow>=10.0.1 # OSRB filed under: https://nvbugs/4702475
requests # OSRB filed under: https://nvbugs/4702237
torch # OSRB filed under: https://nvbugs/4702252
OpenEXR # OSRB filed under: https://nvbugs/5642859