Add 3D infrence support
This commit is contained in:
Sirshendu Ganguly
@@ -0,0 +1,178 @@
|
||||
from .utils import runwareUtils as rwUtils
|
||||
import comfy.model_management
|
||||
|
||||
|
||||
class threeDInference:
|
||||
"""Runware 3D Inference node for generating 3D models from images"""
|
||||
|
||||
# Available 3D models
|
||||
THREED_MODELS = {
|
||||
"Meta SAM 3D": "meta:sam@3d",
|
||||
}
|
||||
|
||||
# Output formats
|
||||
OUTPUT_FORMATS = ["GLB", "PLY"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"Model": (list(cls.THREED_MODELS.keys()), {
|
||||
"tooltip": "Select the 3D generation model to use.",
|
||||
"default": "Meta SAM 3D",
|
||||
}),
|
||||
"positivePrompt": ("STRING", {
|
||||
"multiline": True,
|
||||
"placeholder": "Describe the 3D object you want to generate...",
|
||||
"tooltip": "A text prompt describing the 3D object to generate."
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"tooltip": "Seed for reproducibility. Use the same seed to get the same result.",
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 2147483647,
|
||||
}),
|
||||
"outputFormat": (cls.OUTPUT_FORMATS, {
|
||||
"tooltip": "Output format for the 3D model file.",
|
||||
"default": "GLB",
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"inputs": ("RUNWARE3DINFERENCEINPUTS", {
|
||||
"tooltip": "Connect a Runware 3D Inference Inputs node to provide image and mask inputs.",
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
DESCRIPTION = "Generate 3D models from images using Runware's 3D Inference API. Connect the output to 'Runware Save 3D' node to save the file."
|
||||
FUNCTION = "generate3D"
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("3dObject",)
|
||||
CATEGORY = "Runware"
|
||||
|
||||
def generate3D(self, **kwargs):
|
||||
"""Generate 3D model from inputs"""
|
||||
modelName = kwargs.get("Model", "Meta SAM 3D")
|
||||
positivePrompt = kwargs.get("positivePrompt", "")
|
||||
seed = kwargs.get("seed", 1)
|
||||
outputFormat = kwargs.get("outputFormat", "GLB")
|
||||
inputs = kwargs.get("inputs", None)
|
||||
|
||||
# Get model AIR code
|
||||
model = self.THREED_MODELS.get(modelName, "meta:sam@3d")
|
||||
|
||||
# Build generation config
|
||||
genConfig = [
|
||||
{
|
||||
"taskType": "3dInference",
|
||||
"taskUUID": rwUtils.genRandUUID(),
|
||||
"model": model,
|
||||
"numberResults": 1,
|
||||
"outputType": "URL",
|
||||
"includeCost": True,
|
||||
"deliveryMethod": "async",
|
||||
"outputFormat": outputFormat,
|
||||
}
|
||||
]
|
||||
|
||||
# Add positivePrompt if provided
|
||||
if positivePrompt and positivePrompt.strip() != "":
|
||||
genConfig[0]["positivePrompt"] = positivePrompt.strip()
|
||||
|
||||
# Add seed
|
||||
genConfig[0]["seed"] = seed
|
||||
|
||||
# Handle inputs from 3D Inference Inputs node
|
||||
if inputs is not None and isinstance(inputs, dict) and len(inputs) > 0:
|
||||
genConfig[0]["inputs"] = inputs
|
||||
|
||||
try:
|
||||
print(f"[DEBUG] Sending 3D Inference Request:")
|
||||
print(f"[DEBUG] Request Payload: {rwUtils.safe_json_dumps(genConfig, indent=2)}")
|
||||
|
||||
# Send initial request
|
||||
genResult = rwUtils.inferenecRequest(genConfig)
|
||||
|
||||
print(f"[DEBUG] Received 3D Inference Response:")
|
||||
print(f"[DEBUG] Response: {rwUtils.safe_json_dumps(genResult, indent=2)}")
|
||||
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
# Extract task UUID for polling
|
||||
taskUUID = genConfig[0]["taskUUID"]
|
||||
|
||||
# Poll for 3D generation completion
|
||||
while True:
|
||||
# Check for interrupt before each poll
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
|
||||
# Poll for result
|
||||
pollResult = rwUtils.pollVideoResult(taskUUID) # Reuse existing poll function
|
||||
print(f"[Debugging] Poll result: {rwUtils.safe_json_dumps(pollResult, indent=2) if isinstance(pollResult, (dict, list)) else pollResult}")
|
||||
|
||||
# Check for errors
|
||||
if pollResult and "errors" in pollResult and len(pollResult["errors"]) > 0:
|
||||
error_info = pollResult["errors"][0]
|
||||
error_message = error_info.get("message", "Unknown error")
|
||||
|
||||
if "responseContent" in error_info:
|
||||
response_content = error_info["responseContent"]
|
||||
if isinstance(response_content, str):
|
||||
detailed_message = response_content
|
||||
elif isinstance(response_content, dict):
|
||||
detailed_message = response_content.get("message", str(response_content))
|
||||
else:
|
||||
detailed_message = str(response_content)
|
||||
|
||||
if detailed_message:
|
||||
error_message = f"{error_message}\nProvider Error: {detailed_message}"
|
||||
|
||||
task_uuid = error_info.get("taskUUID", "unknown")
|
||||
raise Exception(f"3D generation failed (Task: {task_uuid}): {error_message}")
|
||||
|
||||
# Check for successful completion
|
||||
if pollResult and "data" in pollResult and len(pollResult["data"]) > 0:
|
||||
data = pollResult["data"][0]
|
||||
|
||||
# Check if outputs.files exists (3D inference response format)
|
||||
outputs = data.get("outputs", {})
|
||||
files = outputs.get("files", [])
|
||||
|
||||
if len(files) > 0:
|
||||
file_info = files[0]
|
||||
file_url = file_info.get("url", "")
|
||||
|
||||
if file_url:
|
||||
print(f"[3D Inference] Generated 3D file URL: {file_url}")
|
||||
return (file_url,)
|
||||
|
||||
# Check status if available
|
||||
if "status" in data:
|
||||
status = data["status"]
|
||||
if status == "success":
|
||||
# Try alternative response format
|
||||
if "outputs" in data and "files" in data["outputs"]:
|
||||
files = data["outputs"]["files"]
|
||||
if len(files) > 0:
|
||||
file_url = files[0].get("url", "")
|
||||
if file_url:
|
||||
print(f"[3D Inference] Generated 3D file URL: {file_url}")
|
||||
return (file_url,)
|
||||
|
||||
# Check for interrupt before waiting
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
|
||||
# Wait before next poll
|
||||
for _ in range(10): # 10 x 0.1 second = 1 second total
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
rwUtils.time.sleep(0.1)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Runware3DInference": threeDInference,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Runware3DInference": "Runware 3D Inference",
|
||||
}
|
||||
Reference in New Issue
Block a user