feat: implement storage endpoints migration for ComfyUI nodes (#8)
feat: migrate ComfyUI nodes from presigned URLs to Brain API storage endpoints ## Core Architecture Changes - Add new BrainApiClient for direct Brain API communication - Replace presigned URL workflow with storage_id/filename pattern - Update MediaStreamInput to use storage_id + filename instead of presigned_download_url - Update MediaStreamOutput to use storage_id + filename instead of presigned_upload_url ## Brain API Client Implementation - Add comprehensive BrainApiClient with authentication and error handling - Implement get_presigned_upload_url() and get_presigned_download_url() methods - Add proper Bearer token authentication for Brain API requests - Include environment variable configuration for Brain API endpoint and credentials ## Media Stream Node Updates - Refactor MediaStreamInput to request presigned URLs from Brain API using storage_id - Update MediaStreamOutput to upload directly to MinIO using Brain API presigned URLs - Simplify node interface by removing presigned URL inputs from workflow templates - Update JavaScript frontend to hide system inputs and manage node UI state ## Infrastructure Improvements - Add .idea/ to .gitignore for IDE file exclusions - Update user_input.py to use new storage_id pattern - Streamline web interface JavaScript for cleaner node management
This commit is contained in:
+1
-1
@@ -157,4 +157,4 @@ cython_debug/
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
.idea/
|
||||
|
||||
@@ -0,0 +1,372 @@
|
||||
"""
|
||||
Brain API Client for ComfyUI Nodes
|
||||
|
||||
This client provides methods to interact with the Brain API storage endpoints,
|
||||
replacing the need for pre-signed URLs in the ComfyUI workflow.
|
||||
"""
|
||||
|
||||
import requests
|
||||
import os
|
||||
import logging
|
||||
from typing import Optional, Dict, Any
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
dotenv_path = os.path.join(current_dir, ".env")
|
||||
load_dotenv(dotenv_path=dotenv_path)
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
|
||||
|
||||
class BrainApiClient:
|
||||
"""
|
||||
Client for interacting with Brain API storage endpoints.
|
||||
|
||||
This client handles authentication and provides methods for uploading,
|
||||
downloading, and deleting files through the Brain API storage endpoints.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the Brain API client with configuration from environment variables."""
|
||||
self.base_url = os.getenv("BRANDO_BRAIN_API_BASE_URL", "http://localhost:2024/api")
|
||||
self.api_key = os.getenv("BRANDO_API_KEY")
|
||||
|
||||
if not self.api_key:
|
||||
raise ValueError(
|
||||
"BRANDO_API_KEY environment variable is required for Brain API authentication"
|
||||
)
|
||||
|
||||
self.headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"User-Agent": "ComfyUI-NilorNodes/1.0"
|
||||
}
|
||||
|
||||
logging.info(f"Brain API Client initialized with base URL: {self.base_url}")
|
||||
|
||||
def upload_file_to_storage(self, file_path: str, filename: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Upload a file to Brain API storage and return storage metadata.
|
||||
|
||||
Args:
|
||||
file_path: Local path to the file to upload
|
||||
filename: Name to use for the uploaded file
|
||||
|
||||
Returns:
|
||||
Dict containing storage_id and filename
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If upload fails
|
||||
FileNotFoundError: If file_path doesn't exist
|
||||
"""
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
url = f"{self.base_url}/storage/upload"
|
||||
|
||||
try:
|
||||
with open(file_path, 'rb') as file:
|
||||
files = {'file': (filename, file, 'application/octet-stream')}
|
||||
|
||||
logging.info(f"Uploading file '{filename}' to Brain API storage...")
|
||||
response = requests.post(
|
||||
url,
|
||||
files=files,
|
||||
headers=self.headers,
|
||||
timeout=300
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
logging.info(f"Upload successful. Storage ID: {result.get('storage_id')}")
|
||||
return result
|
||||
|
||||
except requests.RequestException as e:
|
||||
logging.error(f"Failed to upload file '{filename}': {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Unexpected error uploading file '{filename}': {e}")
|
||||
raise
|
||||
|
||||
def upload_fileobj_to_storage(self, file_obj, filename: str, content_type: str = 'application/octet-stream') -> Dict[str, Any]:
|
||||
"""
|
||||
Upload a file-like object to Brain API storage and return storage metadata.
|
||||
|
||||
Args:
|
||||
file_obj: File-like object to upload
|
||||
filename: Name to use for the uploaded file
|
||||
content_type: MIME type of the file
|
||||
|
||||
Returns:
|
||||
Dict containing storage_id and filename
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If upload fails
|
||||
"""
|
||||
url = f"{self.base_url}/storage/upload"
|
||||
|
||||
try:
|
||||
files = {'file': (filename, file_obj, content_type)}
|
||||
|
||||
logging.info(f"Uploading file object '{filename}' to Brain API storage...")
|
||||
response = requests.post(
|
||||
url,
|
||||
files=files,
|
||||
headers=self.headers,
|
||||
timeout=300
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
logging.info(f"Upload successful. Storage ID: {result.get('storage_id')}")
|
||||
return result
|
||||
|
||||
except requests.RequestException as e:
|
||||
logging.error(f"Failed to upload file object '{filename}': {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Unexpected error uploading file object '{filename}': {e}")
|
||||
raise
|
||||
|
||||
def download_file_from_storage(self, storage_id: str, filename: str, dest_path: str) -> str:
|
||||
"""
|
||||
Download a file from Brain API storage to a local path.
|
||||
|
||||
Args:
|
||||
storage_id: Storage ID of the file to download
|
||||
filename: Name of the file to download
|
||||
dest_path: Local path where the file should be saved
|
||||
|
||||
Returns:
|
||||
Path to the downloaded file
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If download fails
|
||||
"""
|
||||
url = f"{self.base_url}/storage/{storage_id}"
|
||||
params = {'filename': filename}
|
||||
|
||||
try:
|
||||
logging.info(f"Downloading file '{filename}' (storage_id: {storage_id}) from Brain API storage...")
|
||||
response = requests.get(
|
||||
url,
|
||||
params=params,
|
||||
headers=self.headers,
|
||||
timeout=300,
|
||||
stream=True
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
# Ensure destination directory exists
|
||||
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
|
||||
|
||||
with open(dest_path, 'wb') as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
logging.info(f"Download successful. File saved to: {dest_path}")
|
||||
return dest_path
|
||||
|
||||
except requests.RequestException as e:
|
||||
logging.error(f"Failed to download file '{filename}' (storage_id: {storage_id}): {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Unexpected error downloading file '{filename}': {e}")
|
||||
raise
|
||||
|
||||
def get_file_from_storage(self, storage_id: str, filename: str) -> bytes:
|
||||
"""
|
||||
Get file content from Brain API storage as bytes.
|
||||
|
||||
Args:
|
||||
storage_id: Storage ID of the file to download
|
||||
filename: Name of the file to download
|
||||
|
||||
Returns:
|
||||
File content as bytes
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If download fails
|
||||
"""
|
||||
url = f"{self.base_url}/storage/{storage_id}"
|
||||
params = {'filename': filename}
|
||||
|
||||
try:
|
||||
logging.info(f"Getting file '{filename}' (storage_id: {storage_id}) from Brain API storage...")
|
||||
response = requests.get(
|
||||
url,
|
||||
params=params,
|
||||
headers=self.headers,
|
||||
timeout=300
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
logging.info(f"File retrieval successful. Size: {len(response.content)} bytes")
|
||||
return response.content
|
||||
|
||||
except requests.RequestException as e:
|
||||
logging.error(f"Failed to get file '{filename}' (storage_id: {storage_id}): {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Unexpected error getting file '{filename}': {e}")
|
||||
raise
|
||||
|
||||
def delete_file_from_storage(self, storage_id: str, filename: str) -> None:
|
||||
"""
|
||||
Delete a file from Brain API storage.
|
||||
|
||||
Args:
|
||||
storage_id: Storage ID of the file to delete
|
||||
filename: Name of the file to delete
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If deletion fails
|
||||
"""
|
||||
url = f"{self.base_url}/storage/{storage_id}"
|
||||
params = {'filename': filename}
|
||||
|
||||
try:
|
||||
logging.info(f"Deleting file '{filename}' (storage_id: {storage_id}) from Brain API storage...")
|
||||
response = requests.delete(
|
||||
url,
|
||||
params=params,
|
||||
headers=self.headers,
|
||||
timeout=60
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
logging.info(f"File deletion successful")
|
||||
|
||||
except requests.RequestException as e:
|
||||
logging.error(f"Failed to delete file '{filename}' (storage_id: {storage_id}): {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Unexpected error deleting file '{filename}': {e}")
|
||||
raise
|
||||
|
||||
def get_presigned_upload_url(self, filename: str, content_type: str, minio_endpoint: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get a presigned upload URL from Brain API for direct MinIO upload.
|
||||
|
||||
Args:
|
||||
filename: Name of the file to upload
|
||||
content_type: MIME type of the file
|
||||
minio_endpoint: MinIO endpoint that ComfyUI can access
|
||||
|
||||
Returns:
|
||||
Dict containing storage_id, upload_url, and object_key
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If request fails
|
||||
"""
|
||||
url = f"{self.base_url}/storage/generate-upload-url"
|
||||
payload = {
|
||||
"filename": filename,
|
||||
"content_type": content_type,
|
||||
"minio_endpoint": minio_endpoint
|
||||
}
|
||||
|
||||
try:
|
||||
logging.info(f"Requesting presigned upload URL for '{filename}' from Brain API...")
|
||||
response = requests.post(
|
||||
url,
|
||||
json=payload,
|
||||
headers=self.headers,
|
||||
timeout=30
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
logging.info(f"Presigned upload URL generated. Storage ID: {result.get('storage_id')}")
|
||||
return result
|
||||
|
||||
except requests.RequestException as e:
|
||||
logging.error(f"Failed to get presigned upload URL for '{filename}': {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Unexpected error getting presigned upload URL for '{filename}': {e}")
|
||||
raise
|
||||
|
||||
def get_presigned_download_url(self, storage_id: str, filename: str, minio_endpoint: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get a presigned download URL from Brain API for direct MinIO download.
|
||||
|
||||
Args:
|
||||
storage_id: Storage ID of the file to download
|
||||
filename: Name of the file to download
|
||||
minio_endpoint: MinIO endpoint that ComfyUI can access
|
||||
|
||||
Returns:
|
||||
Dict containing download_url
|
||||
|
||||
Raises:
|
||||
requests.RequestException: If request fails
|
||||
"""
|
||||
url = f"{self.base_url}/storage/generate-download-url"
|
||||
payload = {
|
||||
"storage_id": storage_id,
|
||||
"filename": filename,
|
||||
"minio_endpoint": minio_endpoint
|
||||
}
|
||||
|
||||
try:
|
||||
logging.info(f"Requesting presigned download URL for '{filename}' (storage_id: {storage_id}) from Brain API...")
|
||||
response = requests.post(
|
||||
url,
|
||||
json=payload,
|
||||
headers=self.headers,
|
||||
timeout=30
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
logging.info(f"Presigned download URL generated for '{filename}'")
|
||||
return result
|
||||
|
||||
except requests.RequestException as e:
|
||||
logging.error(f"Failed to get presigned download URL for '{filename}' (storage_id: {storage_id}): {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(f"Unexpected error getting presigned download URL for '{filename}': {e}")
|
||||
raise
|
||||
|
||||
def health_check(self) -> bool:
|
||||
"""
|
||||
Check if the Brain API is accessible and authentication is working.
|
||||
|
||||
Returns:
|
||||
True if API is accessible, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Try to access a simple endpoint to verify connectivity
|
||||
url = f"{self.base_url}/health" # Assuming there's a health endpoint
|
||||
response = requests.get(url, headers=self.headers, timeout=10)
|
||||
return response.status_code == 200
|
||||
except:
|
||||
# If health endpoint doesn't exist, try the storage upload endpoint
|
||||
# with a HEAD request to check authentication
|
||||
try:
|
||||
url = f"{self.base_url}/storage/upload"
|
||||
response = requests.head(url, headers=self.headers, timeout=10)
|
||||
return response.status_code in [200, 405] # 405 Method Not Allowed is OK for HEAD
|
||||
except:
|
||||
return False
|
||||
|
||||
|
||||
# Global client instance
|
||||
_brain_api_client = None
|
||||
|
||||
def get_brain_api_client() -> BrainApiClient:
|
||||
"""
|
||||
Get or create the global Brain API client instance.
|
||||
|
||||
Returns:
|
||||
BrainApiClient instance
|
||||
"""
|
||||
global _brain_api_client
|
||||
if _brain_api_client is None:
|
||||
_brain_api_client = BrainApiClient()
|
||||
return _brain_api_client
|
||||
+110
-80
@@ -10,6 +10,7 @@ import boto3
|
||||
import os
|
||||
import json
|
||||
from dotenv import load_dotenv
|
||||
from .brain_api_client import get_brain_api_client
|
||||
|
||||
# --- Load Environment Variables ---
|
||||
# Get the directory of the current script
|
||||
@@ -49,9 +50,13 @@ class MediaStreamInput:
|
||||
{"default": "default_input", "multiline": False},
|
||||
),
|
||||
"format": (["image", "image_batch", "video"],),
|
||||
"presigned_download_url": (
|
||||
"storage_id": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "<auto-filled by system>"},
|
||||
{"default": "<auto-filled by system>", "multiline": False},
|
||||
),
|
||||
"filename": (
|
||||
"STRING",
|
||||
{"default": "<auto-filled by system>", "multiline": False},
|
||||
),
|
||||
},
|
||||
"hidden": {},
|
||||
@@ -64,19 +69,31 @@ class MediaStreamInput:
|
||||
|
||||
def download(
|
||||
self,
|
||||
presigned_download_url: str,
|
||||
storage_id: str,
|
||||
filename: str,
|
||||
format: str,
|
||||
input_name: str = "default_input",
|
||||
):
|
||||
logging.info(
|
||||
f"ℹ️\u2009 Nilor-Nodes: MediaStreamInput: Downloading from {presigned_download_url} for input '{input_name}' with format '{format}'"
|
||||
f"ℹ️\u2009 Nilor-Nodes: MediaStreamInput: Downloading file '{filename}' (storage_id: {storage_id}) for input '{input_name}' with format '{format}'"
|
||||
)
|
||||
try:
|
||||
# Get Brain API client and MinIO endpoint
|
||||
brain_client = get_brain_api_client()
|
||||
minio_endpoint = os.getenv("MINIO_ENDPOINT")
|
||||
if not minio_endpoint:
|
||||
raise ValueError("MINIO_ENDPOINT environment variable is required but not set")
|
||||
|
||||
# Two-phase download for batches: manifest first, then assets
|
||||
if format == "image_batch":
|
||||
manifest_response = requests.get(presigned_download_url, timeout=60)
|
||||
# Get presigned download URL for manifest
|
||||
manifest_url_response = brain_client.get_presigned_download_url(storage_id, filename, minio_endpoint)
|
||||
manifest_url = manifest_url_response["download_url"]
|
||||
|
||||
# Download manifest file directly from MinIO
|
||||
manifest_response = requests.get(manifest_url, timeout=300)
|
||||
manifest_response.raise_for_status()
|
||||
manifest = manifest_response.json()
|
||||
manifest = json.loads(manifest_response.content.decode('utf-8'))
|
||||
|
||||
logging.info(
|
||||
f"ℹ️\u2009 Nilor-Nodes: Processing manifest for '{manifest.get('input_name')}' with {len(manifest.get('files', []))} assets."
|
||||
@@ -87,14 +104,24 @@ class MediaStreamInput:
|
||||
manifest.get("files", []), key=lambda x: x.get("sequence", 0)
|
||||
)
|
||||
|
||||
# Download all assets in parallel
|
||||
# Download all assets using presigned URLs
|
||||
asset_responses = []
|
||||
for file_info in sorted_files:
|
||||
try:
|
||||
resp = requests.get(file_info["presigned_url"], timeout=180)
|
||||
resp.raise_for_status()
|
||||
asset_responses.append(resp.content)
|
||||
except requests.RequestException as e:
|
||||
file_storage_id = file_info.get("storage_id")
|
||||
file_filename = file_info.get("filename")
|
||||
if not file_storage_id or not file_filename:
|
||||
raise ValueError(f"Missing storage_id or filename in manifest file info: {file_info}")
|
||||
|
||||
# Get presigned download URL for this asset
|
||||
asset_url_response = brain_client.get_presigned_download_url(file_storage_id, file_filename, minio_endpoint)
|
||||
asset_url = asset_url_response["download_url"]
|
||||
|
||||
# Download asset directly from MinIO
|
||||
asset_response = requests.get(asset_url, timeout=300)
|
||||
asset_response.raise_for_status()
|
||||
asset_responses.append(asset_response.content)
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
f"🛑\u2009 Nilor-Nodes: Failed to download asset {file_info.get('filename')}: {e}"
|
||||
)
|
||||
@@ -103,9 +130,14 @@ class MediaStreamInput:
|
||||
return self._process_image_batch(asset_responses)
|
||||
|
||||
# --- Single-file download ---
|
||||
response = requests.get(presigned_download_url, timeout=180)
|
||||
response.raise_for_status()
|
||||
media_bytes = response.content
|
||||
# Get presigned download URL
|
||||
download_url_response = brain_client.get_presigned_download_url(storage_id, filename, minio_endpoint)
|
||||
download_url = download_url_response["download_url"]
|
||||
|
||||
# Download file directly from MinIO
|
||||
media_response = requests.get(download_url, timeout=300)
|
||||
media_response.raise_for_status()
|
||||
media_bytes = media_response.content
|
||||
|
||||
if format == "video":
|
||||
return self._process_video(media_bytes)
|
||||
@@ -117,14 +149,9 @@ class MediaStreamInput:
|
||||
f"[🛑] Nilor-Nodes (MediaStreamInput): Unsupported format '{format}' for single media download."
|
||||
)
|
||||
|
||||
except requests.RequestException as e:
|
||||
logging.error(
|
||||
f"🛑\u2009 Nilor-Nodes (MediaStreamInput): Failed to download file: {e}"
|
||||
)
|
||||
return (None,)
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
f"🛑\u2009 Nilor-Nodes (MediaStreamInput): Failed to process media: {e}"
|
||||
f"🛑\u2009 Nilor-Nodes (MediaStreamInput): Failed to download or process media: {e}"
|
||||
)
|
||||
return (None,)
|
||||
|
||||
@@ -219,18 +246,10 @@ class MediaStreamOutput:
|
||||
"STRING",
|
||||
{"default": "<auto-filled by system>", "multiline": False},
|
||||
),
|
||||
"presigned_upload_url": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "<auto-filled by system>"},
|
||||
),
|
||||
"job_completions_queue_url": (
|
||||
"STRING",
|
||||
{"multiline": True, "default": "<auto-filled by system>"},
|
||||
),
|
||||
"output_object_keys": (
|
||||
"STRING",
|
||||
{"multiline": False, "default": "<auto-filled by system>"},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT",
|
||||
@@ -238,8 +257,7 @@ class MediaStreamOutput:
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("uploaded_url",)
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "upload_and_notify"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = category + subcategories["streaming"]
|
||||
@@ -252,9 +270,7 @@ class MediaStreamOutput:
|
||||
venue,
|
||||
canvas,
|
||||
scene,
|
||||
presigned_upload_url,
|
||||
job_completions_queue_url,
|
||||
output_object_keys,
|
||||
framerate,
|
||||
output_name: str = "default_output",
|
||||
prompt=None,
|
||||
@@ -265,35 +281,28 @@ class MediaStreamOutput:
|
||||
"[🛑] Nilor-Nodes (MediaStreamOutput): content_id is a required input for MediaStreamOutput."
|
||||
)
|
||||
|
||||
# The `output_object_keys` is received as a string representation of a dictionary.
|
||||
# We must parse it back into a dictionary.
|
||||
final_outputs_dict = {}
|
||||
try:
|
||||
# The string may use single quotes, so we replace them for valid JSON.
|
||||
final_outputs_dict = json.loads(output_object_keys.replace("'", '"'))
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): FATAL -- Could not parse output_object_keys from string: {output_object_keys}. Error: {e}"
|
||||
)
|
||||
final_outputs_dict = {} # Send empty dict on failure.
|
||||
# No longer need to parse output_object_keys since we use storage_ids directly
|
||||
|
||||
# The presigned_upload_url provided to this node is specific to its output_name.
|
||||
# We don't need to re-select it. We just need to perform the upload.
|
||||
# Upload the media using Brain API client
|
||||
brain_client = get_brain_api_client()
|
||||
storage_result = None
|
||||
|
||||
if format == "png":
|
||||
self._upload_image(images[0], presigned_upload_url)
|
||||
storage_result = self._upload_image(images[0], brain_client, output_name)
|
||||
elif format == "mp4":
|
||||
self._upload_video(images, presigned_upload_url, framerate)
|
||||
storage_result = self._upload_video(images, brain_client, framerate, output_name)
|
||||
|
||||
# This node is responsible for a single output. We find its corresponding object key.
|
||||
output_key_for_this_node = final_outputs_dict.get(output_name)
|
||||
if not output_key_for_this_node:
|
||||
# Use the storage_id from the upload result for the SQS message
|
||||
if not storage_result:
|
||||
logging.error(
|
||||
f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): FATAL -- Could not find object key for output name '{output_name}' in output_object_keys."
|
||||
f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): FATAL -- Upload failed or no storage_id returned."
|
||||
)
|
||||
# Send an empty dictionary to signal failure.
|
||||
final_outputs_for_sqs = {}
|
||||
else:
|
||||
final_outputs_for_sqs = {output_name: output_key_for_this_node}
|
||||
# Use storage_id directly (it's now a string, not a dict)
|
||||
storage_id = storage_result
|
||||
final_outputs_for_sqs = {output_name: storage_id}
|
||||
|
||||
# After upload, send the filtered dictionary of outputs to the SQS queue.
|
||||
completion_message = {
|
||||
@@ -330,9 +339,9 @@ class MediaStreamOutput:
|
||||
)
|
||||
raise # Re-raise to fail the ComfyUI job
|
||||
|
||||
return {"ui": {"images": []}, "result": (presigned_upload_url,)}
|
||||
return {"ui": {"images": []}}
|
||||
|
||||
def _upload_image(self, image_tensor, url):
|
||||
def _upload_image(self, image_tensor, brain_client, output_name):
|
||||
logging.info(
|
||||
"ℹ️\u2009 Nilor-Nodes (MediaStreamOutput): Uploading as PNG image..."
|
||||
)
|
||||
@@ -343,9 +352,30 @@ class MediaStreamOutput:
|
||||
img_pil.save(buffer, format="PNG", compress_level=4)
|
||||
buffer.seek(0)
|
||||
|
||||
self._perform_upload(buffer, url, "image/png")
|
||||
filename = f"{output_name}.png"
|
||||
minio_endpoint = os.getenv("MINIO_ENDPOINT")
|
||||
if not minio_endpoint:
|
||||
raise ValueError("MINIO_ENDPOINT environment variable is required but not set")
|
||||
|
||||
# Get presigned upload URL
|
||||
upload_url_response = brain_client.get_presigned_upload_url(filename, "image/png", minio_endpoint)
|
||||
upload_url = upload_url_response["upload_url"]
|
||||
storage_id = upload_url_response["storage_id"]
|
||||
|
||||
# Upload directly to MinIO
|
||||
buffer.seek(0)
|
||||
upload_response = requests.put(
|
||||
upload_url,
|
||||
data=buffer.getvalue(),
|
||||
headers={"Content-Type": "image/png"},
|
||||
timeout=300
|
||||
)
|
||||
upload_response.raise_for_status()
|
||||
|
||||
logging.info(f"✅ Nilor-Nodes (MediaStreamOutput): PNG image uploaded successfully. Storage ID: {storage_id}")
|
||||
return storage_id
|
||||
|
||||
def _upload_video(self, image_batch_tensor, url, framerate):
|
||||
def _upload_video(self, image_batch_tensor, brain_client, framerate, output_name):
|
||||
logging.info(
|
||||
f"ℹ️\u2009 Nilor-Nodes (MediaStreamOutput): Uploading as MP4 video. Frame count: {len(image_batch_tensor)}"
|
||||
)
|
||||
@@ -359,29 +389,29 @@ class MediaStreamOutput:
|
||||
imageio.mimwrite(buffer, frames, format="mp4", fps=framerate, quality=8)
|
||||
buffer.seek(0)
|
||||
|
||||
self._perform_upload(buffer, url, "video/mp4")
|
||||
filename = f"{output_name}.mp4"
|
||||
minio_endpoint = os.getenv("MINIO_ENDPOINT")
|
||||
if not minio_endpoint:
|
||||
raise ValueError("MINIO_ENDPOINT environment variable is required but not set")
|
||||
|
||||
# Get presigned upload URL
|
||||
upload_url_response = brain_client.get_presigned_upload_url(filename, "video/mp4", minio_endpoint)
|
||||
upload_url = upload_url_response["upload_url"]
|
||||
storage_id = upload_url_response["storage_id"]
|
||||
|
||||
# Upload directly to MinIO
|
||||
buffer.seek(0)
|
||||
upload_response = requests.put(
|
||||
upload_url,
|
||||
data=buffer.getvalue(),
|
||||
headers={"Content-Type": "video/mp4"},
|
||||
timeout=300
|
||||
)
|
||||
upload_response.raise_for_status()
|
||||
|
||||
logging.info(f"✅ Nilor-Nodes (MediaStreamOutput): MP4 video uploaded successfully. Storage ID: {storage_id}")
|
||||
return storage_id
|
||||
|
||||
def _perform_upload(self, buffer, url, content_type):
|
||||
try:
|
||||
logging.info(
|
||||
f"ℹ️\u2009 Nilor-Nodes (MediaStreamOutput): Uploading to {url} with Content-Type: {content_type}"
|
||||
)
|
||||
headers = {"Content-Type": content_type}
|
||||
response = requests.put(
|
||||
url, data=buffer.read(), headers=headers, timeout=300
|
||||
)
|
||||
response.raise_for_status()
|
||||
logging.info("✅ Nilor-Nodes (MediaStreamOutput): Upload successful.")
|
||||
except requests.RequestException as e:
|
||||
logging.error(
|
||||
f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): Failed to upload media: {e}"
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
f"🛑\u2009 Nilor-Nodes (MediaStreamOutput): Failed to process and upload media: {e}"
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
# --- Node Mappings ---
|
||||
@@ -391,6 +421,6 @@ NODE_CLASS_MAPPINGS = {
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MediaStreamInput": "👺 Media Stream Input (URL)",
|
||||
"MediaStreamOutput": "👺 Media Stream Output (URL)",
|
||||
"MediaStreamInput": "👺 Media Stream Input (Storage)",
|
||||
"MediaStreamOutput": "👺 Media Stream Output (Storage)",
|
||||
}
|
||||
|
||||
+1
-1
@@ -59,7 +59,7 @@ class NilorUserInput_Float:
|
||||
"STRING",
|
||||
{"default": "my_float_input", "multiline": False},
|
||||
),
|
||||
"value": ("FLOAT", {"default": 0.0, "step": 0.001}),
|
||||
"value": ("FLOAT", {"default": 0.0}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+37
-87
@@ -19,95 +19,45 @@ function hideWidgets(node, widgetNames) {
|
||||
});
|
||||
}
|
||||
|
||||
function setupMediaStreamOutput(node) {
|
||||
// Hide system inputs by default
|
||||
hideWidgets(node, [
|
||||
"content_id",
|
||||
"venue",
|
||||
"canvas",
|
||||
"scene",
|
||||
"presigned_upload_url",
|
||||
"job_completions_queue_url",
|
||||
"output_object_keys",
|
||||
]);
|
||||
|
||||
const formatWidget = node.widgets.find((w) => w.name === "format");
|
||||
if (!formatWidget) return;
|
||||
|
||||
// Apply current value
|
||||
toggleFramerateWidget(node, formatWidget.value === "mp4");
|
||||
try {
|
||||
const size = node.computeSize();
|
||||
node.onResize?.(size);
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
} catch (_) {}
|
||||
|
||||
// Chain the widget callback once
|
||||
if (!formatWidget.__nilorPatched) {
|
||||
const originalCallback = formatWidget.callback;
|
||||
formatWidget.callback = function (value) {
|
||||
toggleFramerateWidget(node, value === "mp4");
|
||||
try {
|
||||
const size = node.computeSize();
|
||||
node.onResize?.(size);
|
||||
app.graph?.setDirtyCanvas(true, true);
|
||||
} catch (_) {}
|
||||
if (originalCallback) return originalCallback.apply(this, arguments);
|
||||
};
|
||||
formatWidget.__nilorPatched = true;
|
||||
}
|
||||
}
|
||||
|
||||
function setupMediaStreamInput(node) {
|
||||
hideWidgets(node, ["presigned_download_url"]);
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "comfy.nilor-nodes.mediaStream",
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, appInstance) {
|
||||
if (nodeData?.name === "MediaStreamOutput") {
|
||||
const origAdded = nodeType.prototype.onAdded;
|
||||
nodeType.prototype.onAdded = function () {
|
||||
if (typeof origAdded === "function") origAdded.apply(this, arguments);
|
||||
try { setTimeout(() => setupMediaStreamOutput(this), 0); } catch (_) {}
|
||||
};
|
||||
|
||||
const origConfigure = nodeType.prototype.onConfigure;
|
||||
nodeType.prototype.onConfigure = function () {
|
||||
if (typeof origConfigure === "function") origConfigure.apply(this, arguments);
|
||||
try { setTimeout(() => setupMediaStreamOutput(this), 0); } catch (_) {}
|
||||
};
|
||||
}
|
||||
|
||||
if (nodeData?.name === "MediaStreamInput") {
|
||||
const origAddedIn = nodeType.prototype.onAdded;
|
||||
nodeType.prototype.onAdded = function () {
|
||||
if (typeof origAddedIn === "function") origAddedIn.apply(this, arguments);
|
||||
try { setTimeout(() => setupMediaStreamInput(this), 0); } catch (_) {}
|
||||
};
|
||||
|
||||
const origConfigureIn = nodeType.prototype.onConfigure;
|
||||
nodeType.prototype.onConfigure = function () {
|
||||
if (typeof origConfigureIn === "function") origConfigureIn.apply(this, arguments);
|
||||
try { setTimeout(() => setupMediaStreamInput(this), 0); } catch (_) {}
|
||||
};
|
||||
}
|
||||
},
|
||||
|
||||
afterConfigureGraph(graph) {
|
||||
try {
|
||||
(graph?._nodes || graph?.nodes || []).forEach((n) => {
|
||||
if (n?.comfyClass === "MediaStreamOutput") setupMediaStreamOutput(n);
|
||||
if (n?.comfyClass === "MediaStreamInput") setupMediaStreamInput(n);
|
||||
});
|
||||
} catch (e) {
|
||||
console.warn("nilor-media-stream afterConfigureGraph error", e);
|
||||
}
|
||||
},
|
||||
|
||||
nodeCreated(node) {
|
||||
if (node.comfyClass === "MediaStreamOutput") setupMediaStreamOutput(node);
|
||||
if (node.comfyClass === "MediaStreamInput") setupMediaStreamInput(node);
|
||||
if (node.comfyClass === "MediaStreamOutput") {
|
||||
// Hide system inputs by default
|
||||
hideWidgets(node, [
|
||||
"content_id",
|
||||
"venue",
|
||||
"canvas",
|
||||
"scene",
|
||||
"presigned_upload_url",
|
||||
"job_completions_queue_url",
|
||||
"output_object_keys"
|
||||
]);
|
||||
|
||||
const formatWidget = node.widgets.find((w) => w.name === "format");
|
||||
|
||||
// Initial toggle for framerate based on the default format value
|
||||
toggleFramerateWidget(node, formatWidget.value === "mp4");
|
||||
|
||||
// Store original callback to chain it
|
||||
const originalCallback = formatWidget.callback;
|
||||
|
||||
formatWidget.callback = function (value) {
|
||||
toggleFramerateWidget(node, value === "mp4");
|
||||
|
||||
// Recalculate node size after toggling widgets
|
||||
const size = node.computeSize();
|
||||
node.onResize?.(size);
|
||||
|
||||
if (originalCallback) {
|
||||
return originalCallback.apply(this, arguments);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
if (node.comfyClass === "MediaStreamInput") {
|
||||
// Hide system inputs by default
|
||||
hideWidgets(node, ["presigned_download_url"]);
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user