126 lines
4.9 KiB
Python
126 lines
4.9 KiB
Python
from .utils import runwareUtils as rwUtils
|
|
import folder_paths
|
|
import os
|
|
import re
|
|
import requests
|
|
from datetime import datetime
|
|
|
|
|
|
class RunwareSave3D:
|
|
"""Runware Save 3D node for saving 3D model files from URL"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"3dObject": ("STRING", {
|
|
"tooltip": "3D file URL from Runware 3D Inference node."
|
|
}),
|
|
"filename_prefix": ("STRING", {
|
|
"default": "ComfyUI",
|
|
"tooltip": "Prefix for the filename."
|
|
}),
|
|
},
|
|
"optional": {
|
|
"filepath": ("STRING", {
|
|
"placeholder": "Saved file path will appear here after execution.",
|
|
"tooltip": "This field is automatically populated with the saved file path after save.",
|
|
}),
|
|
},
|
|
"hidden": {"node_id": "UNIQUE_ID"},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("filepath",)
|
|
FUNCTION = "save_3d"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Runware"
|
|
DESCRIPTION = "Save 3D model files (GLB/PLY) from Runware 3D Inference to the output folder."
|
|
|
|
def save_3d(self, **kwargs):
|
|
"""Save 3D file from URL"""
|
|
file_url = kwargs.get("3dObject", "")
|
|
filename_prefix = kwargs.get("filename_prefix", "ComfyUI")
|
|
|
|
if not file_url or not file_url.strip():
|
|
raise Exception("No 3D URL provided. Please connect the '3dObject' output from Runware 3D Inference node.")
|
|
|
|
file_url = file_url.strip()
|
|
|
|
# Get output directory
|
|
output_dir = folder_paths.get_output_directory()
|
|
threed_output_dir = os.path.join(output_dir, "3d_models")
|
|
os.makedirs(threed_output_dir, exist_ok=True)
|
|
|
|
# Get extension from URL (handle query parameters)
|
|
url_without_params = file_url.split('?')[0]
|
|
extension = url_without_params.split('.')[-1].lower()
|
|
|
|
if extension not in ["glb", "ply"]:
|
|
raise Exception(f"Unsupported 3D file extension '{extension}' from URL. Expected .glb or .ply.")
|
|
|
|
# Sanitize filename_prefix to prevent path traversal or arbitrary subdir creation
|
|
prefix_basename = os.path.basename(
|
|
filename_prefix.replace(os.sep, "_").replace(os.altsep or os.sep, "_")
|
|
)
|
|
safe_prefix = re.sub(r"[^\w\-]", "", prefix_basename) or "ComfyUI"
|
|
|
|
# Generate filename with timestamp
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
filename = f"{safe_prefix}_{timestamp}.{extension}"
|
|
filepath = os.path.join(threed_output_dir, filename)
|
|
|
|
# Download with retry logic
|
|
max_retries = 10
|
|
retry_delays = [2, 5, 10, 15, 20]
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
response = requests.get(file_url, stream=True, timeout=60)
|
|
print(f"[Runware Save 3D] Download attempt {attempt + 1}: {response.status_code}")
|
|
|
|
if response.status_code == 422:
|
|
print(f"[Runware Save 3D] Server still processing, waiting...")
|
|
rwUtils.time.sleep(retry_delays[min(attempt, len(retry_delays) - 1)])
|
|
continue
|
|
|
|
if response.status_code == 502:
|
|
print(f"[Runware Save 3D] Server error (502), retrying...")
|
|
rwUtils.time.sleep(retry_delays[min(attempt, len(retry_delays) - 1)])
|
|
continue
|
|
|
|
response.raise_for_status()
|
|
|
|
with open(filepath, 'wb') as f:
|
|
for chunk in response.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
|
|
print(f"[Runware Save 3D] Successfully saved to: {filepath}")
|
|
|
|
# Send filepath to UI (like media upload mediaUUID)
|
|
rwUtils.sendSave3DFilepath(filepath, kwargs.get("node_id"))
|
|
|
|
return {"result": (filepath,), "ui": {"text": (f"Saved: {filename}",)}}
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"[Runware Save 3D] Attempt {attempt + 1} failed: {e}")
|
|
if attempt < max_retries - 1:
|
|
rwUtils.time.sleep(retry_delays[min(attempt, len(retry_delays) - 1)])
|
|
continue
|
|
else:
|
|
raise Exception(f"Failed to download 3D file after {max_retries} attempts: {e}")
|
|
except Exception as e:
|
|
print(f"[Runware Save 3D] Unexpected error: {e}")
|
|
raise Exception(f"Failed to save 3D file: {e}")
|
|
|
|
raise Exception("Failed to download 3D file: max retries exceeded")
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RunwareSave3D": RunwareSave3D,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RunwareSave3D": "Runware Save 3D",
|
|
}
|