diff --git a/comfy_api_nodes/apis/rodin_api.py b/comfy_api_nodes/apis/rodin_api.py new file mode 100644 index 000000000..b0cf171fa --- /dev/null +++ b/comfy_api_nodes/apis/rodin_api.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from enum import Enum +from typing import Optional, List +from pydantic import BaseModel, Field + + +class Rodin3DGenerateRequest(BaseModel): + seed: int = Field(..., description="seed_") + tier: str = Field(..., description="Tier of generation.") + material: str = Field(..., description="The material type.") + quality: str = Field(..., description="The generation quality of the mesh.") + mesh_mode: str = Field(..., description="It controls the type of faces of generated models.") + +class GenerateJobsData(BaseModel): + uuids: List[str] = Field(..., description="str LIST") + subscription_key: str = Field(..., description="subscription key") + +class Rodin3DGenerateResponse(BaseModel): + message: Optional[str] = Field(None, description="Return message.") + prompt: Optional[str] = Field(None, description="Generated Prompt from image.") + submit_time: Optional[str] = Field(None, description="Submit Time") + uuid: Optional[str] = Field(None, description="Task str") + jobs: Optional[GenerateJobsData] = Field(None, description="Details of jobs") + +class JobStatus(str, Enum): + """ + Status for jobs + """ + Done = "Done" + Failed = "Failed" + Generating = "Generating" + Waiting = "Waiting" + +class Rodin3DCheckStatusRequest(BaseModel): + subscription_key: str = Field(..., description="subscription from generate endpoint") + +class JobItem(BaseModel): + uuid: str = Field(..., description="uuid") + status: JobStatus = Field(...,description="Status Currently") + +class Rodin3DCheckStatusResponse(BaseModel): + jobs: List[JobItem] = Field(..., description="Job status List") + +class Rodin3DDownloadRequest(BaseModel): + task_uuid: str = Field(..., description="Task str") + +class RodinResourceItem(BaseModel): + url: str = Field(..., description="Download Url") + name: str = Field(..., description="File name with ext") + +class Rodin3DDownloadResponse(BaseModel): + list: List[RodinResourceItem] = Field(..., description="Source List") + + + + diff --git a/comfy_api_nodes/nodes_rodin.py b/comfy_api_nodes/nodes_rodin.py new file mode 100644 index 000000000..c8b59bf35 --- /dev/null +++ b/comfy_api_nodes/nodes_rodin.py @@ -0,0 +1,462 @@ +""" +ComfyUI X Rodin3D(Deemos) API Nodes + +Rodin API docs: https://developer.hyper3d.ai/ + +""" + +from __future__ import annotations +from inspect import cleandoc +from comfy.comfy_types.node_typing import IO +import folder_paths as comfy_paths +import requests +import os +import datetime +import shutil +import time +import io +import logging +import math +from PIL import Image +from comfy_api_nodes.apis.rodin_api import ( + Rodin3DGenerateRequest, + Rodin3DGenerateResponse, + Rodin3DCheckStatusRequest, + Rodin3DCheckStatusResponse, + Rodin3DDownloadRequest, + Rodin3DDownloadResponse, + JobStatus, +) +from comfy_api_nodes.apis.client import ( + ApiEndpoint, + HttpMethod, + SynchronousOperation, + PollingOperation, +) + + +COMMON_PARAMETERS = { + "Seed": ( + IO.INT, + { + "default":0, + "min":0, + "max":65535, + "display":"number" + } + ), + "Material_Type": ( + IO.COMBO, + { + "options": ["PBR", "Shaded"], + "default": "PBR" + } + ), + "Polygon_count": ( + IO.COMBO, + { + "options": ["4K-Quad", "8K-Quad", "18K-Quad", "50K-Quad", "200K-Triangle"], + "default": "18K-Quad" + } + ) +} + +def create_task_error(response: Rodin3DGenerateResponse): + """Check if the response has error""" + return hasattr(response, "error") + + + +class Rodin3DAPI: + """ + Generate 3D Assets using Rodin API + """ + RETURN_TYPES = (IO.STRING,) + RETURN_NAMES = ("3D Model Path",) + CATEGORY = "api node/3D/Rodin" + DESCRIPTION = cleandoc(__doc__ or "") + FUNCTION = "api_call" + API_NODE = True + + def tensor_to_filelike(self, 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 = io.BytesIO() + image.save(img_byte_arr, format='PNG') # PNG is used for lossless compression + img_byte_arr.seek(0) + return img_byte_arr + + def check_rodin_status(self, response: Rodin3DCheckStatusResponse) -> str: + has_failed = any(job.status == JobStatus.Failed for job in response.jobs) + all_done = all(job.status == JobStatus.Done for job in response.jobs) + status_list = [str(job.status) for job in response.jobs] + logging.info(f"[ Rodin3D API - CheckStatus ] Generate Status: {status_list}") + if has_failed: + logging.error(f"[ Rodin3D API - CheckStatus ] Generate Failed: {status_list}, Please try again.") + raise Exception("[ Rodin3D API ] Generate Failed, Please Try again.") + elif all_done: + return "DONE" + else: + return "Generating" + + def CreateGenerateTask(self, images=None, seed=1, material="PBR", quality="medium", tier="Regular", mesh_mode="Quad", **kwargs): + if images == 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.") + + path = "/proxy/rodin/api/v2/rodin" + operation = SynchronousOperation( + endpoint=ApiEndpoint( + path=path, + method=HttpMethod.POST, + request_model=Rodin3DGenerateRequest, + response_model=Rodin3DGenerateResponse, + ), + request=Rodin3DGenerateRequest( + seed=seed, + tier=tier, + material=material, + quality=quality, + mesh_mode=mesh_mode + ), + files=[ + ( + "images", + open(image, "rb") if isinstance(image, str) else self.tensor_to_filelike(image) + ) + for image in images if image is not None + ], + content_type = "multipart/form-data", + auth_kwargs=kwargs, + ) + + response = operation.execute() + + if create_task_error(response): + 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(f"[ Rodin3D API - Submit Jobs ] UUID: {task_uuid}") + return task_uuid, subscription_key + + def poll_for_task_status(self, subscription_key, **kwargs) -> Rodin3DCheckStatusResponse: + + path = "/proxy/rodin/api/v2/status" + + poll_operation = PollingOperation( + poll_endpoint=ApiEndpoint( + path = path, + method=HttpMethod.POST, + request_model=Rodin3DCheckStatusRequest, + response_model=Rodin3DCheckStatusResponse, + ), + request=Rodin3DCheckStatusRequest( + subscription_key = subscription_key + ), + completed_statuses=["DONE"], + failed_statuses=["FAILED"], + status_extractor=self.check_rodin_status, + poll_interval=3.0, + auth_kwargs=kwargs, + ) + + logging.info("[ Rodin3D API - CheckStatus ] Generate Start!") + + return poll_operation.execute() + + + + def GetRodinDownloadList(self, uuid, **kwargs) -> Rodin3DDownloadResponse: + logging.info("[ Rodin3D API - Downloading ] Generate Successfully!") + + path = "/proxy/rodin/api/v2/download" + operation = SynchronousOperation( + endpoint=ApiEndpoint( + path=path, + method=HttpMethod.POST, + request_model=Rodin3DDownloadRequest, + response_model=Rodin3DDownloadResponse, + ), + request=Rodin3DDownloadRequest( + task_uuid=uuid + ), + auth_kwargs=kwargs + ) + + return operation.execute() + + def GetQualityAndMode(self, PolyCount): + if PolyCount == "200K-Triangle": + mesh_mode = "Raw" + quality = "medium" + else: + mesh_mode = "Quad" + if PolyCount == "4K-Quad": + quality = "extra-low" + elif PolyCount == "8K-Quad": + quality = "low" + elif PolyCount == "18K-Quad": + quality = "medium" + elif PolyCount == "50K-Quad": + quality = "high" + else: + quality = "medium" + + return mesh_mode, quality + + def DownLoadFiles(self, Url_List): + Save_path = os.path.join(comfy_paths.get_output_directory(), "Rodin3D", datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")) + os.makedirs(Save_path, exist_ok=True) + model_file_path = None + for Item in Url_List.list: + url = Item.url + file_name = Item.name + file_path = os.path.join(Save_path, file_name) + if file_path.endswith(".glb"): + model_file_path = file_path + logging.info(f"[ Rodin3D API - download_files ] Downloading file: {file_path}") + max_retries = 5 + for attempt in range(max_retries): + try: + with requests.get(url, stream=True) as r: + r.raise_for_status() + with open(file_path, "wb") as f: + shutil.copyfileobj(r.raw, f) + break + except Exception as e: + logging.info(f"[ Rodin3D API - download_files ] Error downloading {file_path}:{e}") + if attempt < max_retries - 1: + logging.info("Retrying...") + time.sleep(2) + else: + logging.info(f"[ Rodin3D API - download_files ] Failed to download {file_path} after {max_retries} attempts.") + + return model_file_path + + +class Rodin3D_Regular(Rodin3DAPI): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "Images": + ( + IO.IMAGE, + { + "forceInput":True, + } + ) + }, + "optional": { + **COMMON_PARAMETERS + }, + "hidden": { + "auth_token": "AUTH_TOKEN_COMFY_ORG", + "comfy_api_key": "API_KEY_COMFY_ORG", + }, + } + + def api_call( + self, + Images, + Seed, + Material_Type, + Polygon_count, + **kwargs + ): + tier = "Regular" + num_images = Images.shape[0] + m_images = [] + for i in range(num_images): + m_images.append(Images[i]) + mesh_mode, quality = self.GetQualityAndMode(Polygon_count) + task_uuid, subscription_key = self.CreateGenerateTask(images=m_images, seed=Seed, material=Material_Type, quality=quality, tier=tier, mesh_mode=mesh_mode, **kwargs) + self.poll_for_task_status(subscription_key, **kwargs) + Download_List = self.GetRodinDownloadList(task_uuid, **kwargs) + model = self.DownLoadFiles(Download_List) + + return (model,) + +class Rodin3D_Detail(Rodin3DAPI): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "Images": + ( + IO.IMAGE, + { + "forceInput":True, + } + ) + }, + "optional": { + **COMMON_PARAMETERS + }, + "hidden": { + "auth_token": "AUTH_TOKEN_COMFY_ORG", + "comfy_api_key": "API_KEY_COMFY_ORG", + }, + } + + def api_call( + self, + Images, + Seed, + Material_Type, + Polygon_count, + **kwargs + ): + tier = "Detail" + num_images = Images.shape[0] + m_images = [] + for i in range(num_images): + m_images.append(Images[i]) + mesh_mode, quality = self.GetQualityAndMode(Polygon_count) + task_uuid, subscription_key = self.CreateGenerateTask(images=m_images, seed=Seed, material=Material_Type, quality=quality, tier=tier, mesh_mode=mesh_mode, **kwargs) + self.poll_for_task_status(subscription_key, **kwargs) + Download_List = self.GetRodinDownloadList(task_uuid, **kwargs) + model = self.DownLoadFiles(Download_List) + + return (model,) + +class Rodin3D_Smooth(Rodin3DAPI): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "Images": + ( + IO.IMAGE, + { + "forceInput":True, + } + ) + }, + "optional": { + **COMMON_PARAMETERS + }, + "hidden": { + "auth_token": "AUTH_TOKEN_COMFY_ORG", + "comfy_api_key": "API_KEY_COMFY_ORG", + }, + } + + def api_call( + self, + Images, + Seed, + Material_Type, + Polygon_count, + **kwargs + ): + tier = "Smooth" + num_images = Images.shape[0] + m_images = [] + for i in range(num_images): + m_images.append(Images[i]) + mesh_mode, quality = self.GetQualityAndMode(Polygon_count) + task_uuid, subscription_key = self.CreateGenerateTask(images=m_images, seed=Seed, material=Material_Type, quality=quality, tier=tier, mesh_mode=mesh_mode, **kwargs) + self.poll_for_task_status(subscription_key, **kwargs) + Download_List = self.GetRodinDownloadList(task_uuid, **kwargs) + model = self.DownLoadFiles(Download_List) + + return (model,) + +class Rodin3D_Sketch(Rodin3DAPI): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "Images": + ( + IO.IMAGE, + { + "forceInput":True, + } + ) + }, + "optional": { + "Seed": + ( + IO.INT, + { + "default":0, + "min":0, + "max":65535, + "display":"number" + } + ) + }, + "hidden": { + "auth_token": "AUTH_TOKEN_COMFY_ORG", + "comfy_api_key": "API_KEY_COMFY_ORG", + }, + } + + def api_call( + self, + Images, + Seed, + **kwargs + ): + tier = "Sketch" + num_images = Images.shape[0] + m_images = [] + for i in range(num_images): + m_images.append(Images[i]) + material_type = "PBR" + quality = "medium" + mesh_mode = "Quad" + task_uuid, subscription_key = self.CreateGenerateTask(images=m_images, seed=Seed, material=material_type, quality=quality, tier=tier, mesh_mode=mesh_mode, **kwargs) + self.poll_for_task_status(subscription_key, **kwargs) + Download_List = self.GetRodinDownloadList(task_uuid, **kwargs) + model = self.DownLoadFiles(Download_List) + + return (model,) + +# A dictionary that contains all nodes you want to export with their names +# NOTE: names should be globally unique +NODE_CLASS_MAPPINGS = { + "Rodin3D_Regular": Rodin3D_Regular, + "Rodin3D_Detail": Rodin3D_Detail, + "Rodin3D_Smooth": Rodin3D_Smooth, + "Rodin3D_Sketch": Rodin3D_Sketch, +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "Rodin3D_Regular": "Rodin 3D Generate - Regular Generate", + "Rodin3D_Detail": "Rodin 3D Generate - Detail Generate", + "Rodin3D_Smooth": "Rodin 3D Generate - Smooth Generate", + "Rodin3D_Sketch": "Rodin 3D Generate - Sketch Generate", +} diff --git a/nodes.py b/nodes.py index 4bb54e7d5..772e95fac 100644 --- a/nodes.py +++ b/nodes.py @@ -2282,6 +2282,7 @@ def init_builtin_api_nodes(): "nodes_stability.py", "nodes_pika.py", "nodes_runway.py", + "nodes_rodin.py", "nodes_gemini.py", ]