mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2025-12-08 21:44:33 +08:00
521 lines
16 KiB
Python
521 lines
16 KiB
Python
"""
|
|
ComfyUI X Rodin3D(Deemos) API Nodes
|
|
|
|
Rodin API docs: https://developer.hyper3d.ai/
|
|
|
|
"""
|
|
|
|
from inspect import cleandoc
|
|
import folder_paths as comfy_paths
|
|
import os
|
|
import logging
|
|
import math
|
|
from typing import Optional
|
|
from io import BytesIO
|
|
from typing_extensions import override
|
|
from PIL import Image
|
|
from comfy_api_nodes.apis.rodin_api import (
|
|
Rodin3DGenerateRequest,
|
|
Rodin3DGenerateResponse,
|
|
Rodin3DCheckStatusRequest,
|
|
Rodin3DCheckStatusResponse,
|
|
Rodin3DDownloadRequest,
|
|
Rodin3DDownloadResponse,
|
|
JobStatus,
|
|
)
|
|
from comfy_api_nodes.util import (
|
|
sync_op,
|
|
poll_op,
|
|
ApiEndpoint,
|
|
download_url_to_bytesio,
|
|
)
|
|
from comfy_api.latest import ComfyExtension, IO
|
|
|
|
|
|
COMMON_PARAMETERS = [
|
|
IO.Int.Input(
|
|
"Seed",
|
|
default=0,
|
|
min=0,
|
|
max=65535,
|
|
display_mode=IO.NumberDisplay.number,
|
|
optional=True,
|
|
),
|
|
IO.Combo.Input("Material_Type", options=["PBR", "Shaded"], default="PBR", optional=True),
|
|
IO.Combo.Input(
|
|
"Polygon_count",
|
|
options=["4K-Quad", "8K-Quad", "18K-Quad", "50K-Quad", "200K-Triangle"],
|
|
default="18K-Quad",
|
|
optional=True,
|
|
),
|
|
]
|
|
|
|
|
|
def get_quality_mode(poly_count):
|
|
polycount = poly_count.split("-")
|
|
poly = polycount[1]
|
|
count = polycount[0]
|
|
if poly == "Triangle":
|
|
mesh_mode = "Raw"
|
|
elif poly == "Quad":
|
|
mesh_mode = "Quad"
|
|
else:
|
|
mesh_mode = "Quad"
|
|
|
|
if count == "4K":
|
|
quality_override = 4000
|
|
elif count == "8K":
|
|
quality_override = 8000
|
|
elif count == "18K":
|
|
quality_override = 18000
|
|
elif count == "50K":
|
|
quality_override = 50000
|
|
elif count == "2K":
|
|
quality_override = 2000
|
|
elif count == "20K":
|
|
quality_override = 20000
|
|
elif count == "150K":
|
|
quality_override = 150000
|
|
elif count == "500K":
|
|
quality_override = 500000
|
|
else:
|
|
quality_override = 18000
|
|
|
|
return mesh_mode, quality_override
|
|
|
|
|
|
def tensor_to_filelike(tensor, max_pixels: int = 2048*2048):
|
|
"""
|
|
Converts a PyTorch tensor to a file-like object.
|
|
|
|
Args:
|
|
- tensor (torch.Tensor): A tensor representing an image of shape (H, W, C)
|
|
where C is the number of channels (3 for RGB), H is height, and W is width.
|
|
|
|
Returns:
|
|
- io.BytesIO: A file-like object containing the image data.
|
|
"""
|
|
array = tensor.cpu().numpy()
|
|
array = (array * 255).astype('uint8')
|
|
image = Image.fromarray(array, 'RGB')
|
|
|
|
original_width, original_height = image.size
|
|
original_pixels = original_width * original_height
|
|
if original_pixels > max_pixels:
|
|
scale = math.sqrt(max_pixels / original_pixels)
|
|
new_width = int(original_width * scale)
|
|
new_height = int(original_height * scale)
|
|
else:
|
|
new_width, new_height = original_width, original_height
|
|
|
|
if new_width != original_width or new_height != original_height:
|
|
image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
|
|
img_byte_arr = BytesIO()
|
|
image.save(img_byte_arr, format='PNG') # PNG is used for lossless compression
|
|
img_byte_arr.seek(0)
|
|
return img_byte_arr
|
|
|
|
|
|
async def create_generate_task(
|
|
cls: type[IO.ComfyNode],
|
|
images=None,
|
|
seed=1,
|
|
material="PBR",
|
|
quality_override=18000,
|
|
tier="Regular",
|
|
mesh_mode="Quad",
|
|
ta_pose: bool = False,
|
|
):
|
|
if images is None:
|
|
raise Exception("Rodin 3D generate requires at least 1 image.")
|
|
if len(images) > 5:
|
|
raise Exception("Rodin 3D generate requires up to 5 image.")
|
|
|
|
response = await sync_op(
|
|
cls,
|
|
ApiEndpoint(path="/proxy/rodin/api/v2/rodin", method="POST"),
|
|
response_model=Rodin3DGenerateResponse,
|
|
data=Rodin3DGenerateRequest(
|
|
seed=seed,
|
|
tier=tier,
|
|
material=material,
|
|
quality_override=quality_override,
|
|
mesh_mode=mesh_mode,
|
|
TAPose=ta_pose,
|
|
),
|
|
files=[
|
|
(
|
|
"images",
|
|
open(image, "rb") if isinstance(image, str) else tensor_to_filelike(image)
|
|
)
|
|
for image in images if image is not None
|
|
],
|
|
content_type="multipart/form-data",
|
|
)
|
|
|
|
if hasattr(response, "error"):
|
|
error_message = f"Rodin3D Create 3D generate Task Failed. Message: {response.message}, error: {response.error}"
|
|
logging.error(error_message)
|
|
raise Exception(error_message)
|
|
|
|
logging.info("[ Rodin3D API - Submit Jobs ] Submit Generate Task Success!")
|
|
subscription_key = response.jobs.subscription_key
|
|
task_uuid = response.uuid
|
|
logging.info("[ Rodin3D API - Submit Jobs ] UUID: %s", task_uuid)
|
|
return task_uuid, subscription_key
|
|
|
|
|
|
def check_rodin_status(response: Rodin3DCheckStatusResponse) -> str:
|
|
all_done = all(job.status == JobStatus.Done for job in response.jobs)
|
|
status_list = [str(job.status) for job in response.jobs]
|
|
logging.info("[ Rodin3D API - CheckStatus ] Generate Status: %s", status_list)
|
|
if any(job.status == JobStatus.Failed for job in response.jobs):
|
|
logging.error("[ Rodin3D API - CheckStatus ] Generate Failed: %s, Please try again.", status_list)
|
|
raise Exception("[ Rodin3D API ] Generate Failed, Please Try again.")
|
|
if all_done:
|
|
return "DONE"
|
|
return "Generating"
|
|
|
|
def extract_progress(response: Rodin3DCheckStatusResponse) -> Optional[int]:
|
|
if not response.jobs:
|
|
return None
|
|
completed_count = sum(1 for job in response.jobs if job.status == JobStatus.Done)
|
|
return int((completed_count / len(response.jobs)) * 100)
|
|
|
|
|
|
async def poll_for_task_status(subscription_key: str, cls: type[IO.ComfyNode]) -> Rodin3DCheckStatusResponse:
|
|
logging.info("[ Rodin3D API - CheckStatus ] Generate Start!")
|
|
return await poll_op(
|
|
cls,
|
|
ApiEndpoint(path="/proxy/rodin/api/v2/status", method="POST"),
|
|
response_model=Rodin3DCheckStatusResponse,
|
|
data=Rodin3DCheckStatusRequest(subscription_key=subscription_key),
|
|
status_extractor=check_rodin_status,
|
|
progress_extractor=extract_progress,
|
|
)
|
|
|
|
|
|
async def get_rodin_download_list(uuid: str, cls: type[IO.ComfyNode]) -> Rodin3DDownloadResponse:
|
|
logging.info("[ Rodin3D API - Downloading ] Generate Successfully!")
|
|
return await sync_op(
|
|
cls,
|
|
ApiEndpoint(path="/proxy/rodin/api/v2/download", method="POST"),
|
|
response_model=Rodin3DDownloadResponse,
|
|
data=Rodin3DDownloadRequest(task_uuid=uuid),
|
|
monitor_progress=False,
|
|
)
|
|
|
|
|
|
async def download_files(url_list, task_uuid: str):
|
|
result_folder_name = f"Rodin3D_{task_uuid}"
|
|
save_path = os.path.join(comfy_paths.get_output_directory(), result_folder_name)
|
|
os.makedirs(save_path, exist_ok=True)
|
|
model_file_path = None
|
|
for i in url_list.list:
|
|
file_path = os.path.join(save_path, i.name)
|
|
if file_path.endswith(".glb"):
|
|
model_file_path = os.path.join(result_folder_name, i.name)
|
|
await download_url_to_bytesio(i.url, file_path)
|
|
return model_file_path
|
|
|
|
|
|
class Rodin3D_Regular(IO.ComfyNode):
|
|
"""Generate 3D Assets using Rodin API"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="Rodin3D_Regular",
|
|
display_name="Rodin 3D Generate - Regular Generate",
|
|
category="api node/3d/Rodin",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.Image.Input("Images"),
|
|
*COMMON_PARAMETERS,
|
|
],
|
|
outputs=[IO.String.Output(display_name="3D Model Path")],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
Images,
|
|
Seed,
|
|
Material_Type,
|
|
Polygon_count,
|
|
) -> IO.NodeOutput:
|
|
tier = "Regular"
|
|
num_images = Images.shape[0]
|
|
m_images = []
|
|
for i in range(num_images):
|
|
m_images.append(Images[i])
|
|
mesh_mode, quality_override = get_quality_mode(Polygon_count)
|
|
task_uuid, subscription_key = await create_generate_task(
|
|
cls,
|
|
images=m_images,
|
|
seed=Seed,
|
|
material=Material_Type,
|
|
quality_override=quality_override,
|
|
tier=tier,
|
|
mesh_mode=mesh_mode,
|
|
)
|
|
await poll_for_task_status(subscription_key, cls)
|
|
download_list = await get_rodin_download_list(task_uuid, cls)
|
|
model = await download_files(download_list, task_uuid)
|
|
|
|
return IO.NodeOutput(model)
|
|
|
|
|
|
class Rodin3D_Detail(IO.ComfyNode):
|
|
"""Generate 3D Assets using Rodin API"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="Rodin3D_Detail",
|
|
display_name="Rodin 3D Generate - Detail Generate",
|
|
category="api node/3d/Rodin",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.Image.Input("Images"),
|
|
*COMMON_PARAMETERS,
|
|
],
|
|
outputs=[IO.String.Output(display_name="3D Model Path")],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
Images,
|
|
Seed,
|
|
Material_Type,
|
|
Polygon_count,
|
|
) -> IO.NodeOutput:
|
|
tier = "Detail"
|
|
num_images = Images.shape[0]
|
|
m_images = []
|
|
for i in range(num_images):
|
|
m_images.append(Images[i])
|
|
mesh_mode, quality_override = get_quality_mode(Polygon_count)
|
|
task_uuid, subscription_key = await create_generate_task(
|
|
cls,
|
|
images=m_images,
|
|
seed=Seed,
|
|
material=Material_Type,
|
|
quality_override=quality_override,
|
|
tier=tier,
|
|
mesh_mode=mesh_mode,
|
|
)
|
|
await poll_for_task_status(subscription_key, cls)
|
|
download_list = await get_rodin_download_list(task_uuid, cls)
|
|
model = await download_files(download_list, task_uuid)
|
|
|
|
return IO.NodeOutput(model)
|
|
|
|
|
|
class Rodin3D_Smooth(IO.ComfyNode):
|
|
"""Generate 3D Assets using Rodin API"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="Rodin3D_Smooth",
|
|
display_name="Rodin 3D Generate - Smooth Generate",
|
|
category="api node/3d/Rodin",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.Image.Input("Images"),
|
|
*COMMON_PARAMETERS,
|
|
],
|
|
outputs=[IO.String.Output(display_name="3D Model Path")],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
Images,
|
|
Seed,
|
|
Material_Type,
|
|
Polygon_count,
|
|
) -> IO.NodeOutput:
|
|
num_images = Images.shape[0]
|
|
m_images = []
|
|
for i in range(num_images):
|
|
m_images.append(Images[i])
|
|
mesh_mode, quality_override = get_quality_mode(Polygon_count)
|
|
task_uuid, subscription_key = await create_generate_task(
|
|
cls,
|
|
images=m_images,
|
|
seed=Seed,
|
|
material=Material_Type,
|
|
quality_override=quality_override,
|
|
tier="Smooth",
|
|
mesh_mode=mesh_mode,
|
|
)
|
|
await poll_for_task_status(subscription_key, cls)
|
|
download_list = await get_rodin_download_list(task_uuid, cls)
|
|
model = await download_files(download_list, task_uuid)
|
|
|
|
return IO.NodeOutput(model)
|
|
|
|
|
|
class Rodin3D_Sketch(IO.ComfyNode):
|
|
"""Generate 3D Assets using Rodin API"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="Rodin3D_Sketch",
|
|
display_name="Rodin 3D Generate - Sketch Generate",
|
|
category="api node/3d/Rodin",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.Image.Input("Images"),
|
|
IO.Int.Input(
|
|
"Seed",
|
|
default=0,
|
|
min=0,
|
|
max=65535,
|
|
display_mode=IO.NumberDisplay.number,
|
|
optional=True,
|
|
),
|
|
],
|
|
outputs=[IO.String.Output(display_name="3D Model Path")],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
Images,
|
|
Seed,
|
|
) -> IO.NodeOutput:
|
|
num_images = Images.shape[0]
|
|
m_images = []
|
|
for i in range(num_images):
|
|
m_images.append(Images[i])
|
|
task_uuid, subscription_key = await create_generate_task(
|
|
cls,
|
|
images=m_images,
|
|
seed=Seed,
|
|
material="PBR",
|
|
quality_override=18000,
|
|
tier="Sketch",
|
|
mesh_mode="Quad",
|
|
)
|
|
await poll_for_task_status(subscription_key, cls)
|
|
download_list = await get_rodin_download_list(task_uuid, cls)
|
|
model = await download_files(download_list, task_uuid)
|
|
|
|
return IO.NodeOutput(model)
|
|
|
|
|
|
class Rodin3D_Gen2(IO.ComfyNode):
|
|
"""Generate 3D Assets using Rodin API"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> IO.Schema:
|
|
return IO.Schema(
|
|
node_id="Rodin3D_Gen2",
|
|
display_name="Rodin 3D Generate - Gen-2 Generate",
|
|
category="api node/3d/Rodin",
|
|
description=cleandoc(cls.__doc__ or ""),
|
|
inputs=[
|
|
IO.Image.Input("Images"),
|
|
IO.Int.Input(
|
|
"Seed",
|
|
default=0,
|
|
min=0,
|
|
max=65535,
|
|
display_mode=IO.NumberDisplay.number,
|
|
optional=True,
|
|
),
|
|
IO.Combo.Input("Material_Type", options=["PBR", "Shaded"], default="PBR", optional=True),
|
|
IO.Combo.Input(
|
|
"Polygon_count",
|
|
options=["4K-Quad", "8K-Quad", "18K-Quad", "50K-Quad", "2K-Triangle", "20K-Triangle", "150K-Triangle", "500K-Triangle"],
|
|
default="500K-Triangle",
|
|
optional=True,
|
|
),
|
|
IO.Boolean.Input("TAPose", default=False),
|
|
],
|
|
outputs=[IO.String.Output(display_name="3D Model Path")],
|
|
hidden=[
|
|
IO.Hidden.auth_token_comfy_org,
|
|
IO.Hidden.api_key_comfy_org,
|
|
IO.Hidden.unique_id,
|
|
],
|
|
is_api_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
async def execute(
|
|
cls,
|
|
Images,
|
|
Seed,
|
|
Material_Type,
|
|
Polygon_count,
|
|
TAPose,
|
|
) -> IO.NodeOutput:
|
|
tier = "Gen-2"
|
|
num_images = Images.shape[0]
|
|
m_images = []
|
|
for i in range(num_images):
|
|
m_images.append(Images[i])
|
|
mesh_mode, quality_override = get_quality_mode(Polygon_count)
|
|
task_uuid, subscription_key = await create_generate_task(
|
|
cls,
|
|
images=m_images,
|
|
seed=Seed,
|
|
material=Material_Type,
|
|
quality_override=quality_override,
|
|
tier=tier,
|
|
mesh_mode=mesh_mode,
|
|
ta_pose=TAPose,
|
|
)
|
|
await poll_for_task_status(subscription_key, cls)
|
|
download_list = await get_rodin_download_list(task_uuid, cls)
|
|
model = await download_files(download_list, task_uuid)
|
|
|
|
return IO.NodeOutput(model)
|
|
|
|
|
|
class Rodin3DExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
|
return [
|
|
Rodin3D_Regular,
|
|
Rodin3D_Detail,
|
|
Rodin3D_Smooth,
|
|
Rodin3D_Sketch,
|
|
Rodin3D_Gen2,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> Rodin3DExtension:
|
|
return Rodin3DExtension()
|