110 lines
4.0 KiB
Python
110 lines
4.0 KiB
Python
from .utils import runwareUtils as rwUtils
|
|
import folder_paths
|
|
import os
|
|
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."
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
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()
|
|
|
|
# Validate extension
|
|
if extension not in ["glb"]:
|
|
extension = "glb" # Default to GLB if unknown
|
|
|
|
# Generate filename with timestamp
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
filename = f"{filename_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}")
|
|
|
|
# Return UI info for ComfyUI
|
|
return {"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",
|
|
}
|