nodes/nilor-corp-nilor-nodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: implement storage endpoints migration for ComfyUI nodes

Browse Source 
- Add BrainApiClient for interacting with Brain API storage endpoints
- Update MediaStreamInput to use storage_id + filename instead of presigned URLs
- Update MediaStreamOutput to use Brain API client for uploads and remove output_object_keys dependency
- Add environment configuration for Brain API connection

This enables ComfyUI nodes to work with the new storage architecture
where Brain API acts as a proxy for MinIO operations, resolving network
topology issues in hybrid Docker/local development setups.

(cherry picked from commit a5f6f05908b081b82a64fbd55b1d113b3d826f3a)
This commit is contained in:
Sylvester Meighan
2025-10-02 14:19:13 -07:00
parent bcd62a62b4
commit a754e83efd
2 changed files with 337 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
brain_api_client.py
+286
View File
@@ -0,0 +1,286 @@
"""
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 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
media_stream.py
+51 -81
View File
@@ -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,23 @@ 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
brain_client = get_brain_api_client()
# Two-phase download for batches: manifest first, then assets
if format == "image_batch":
manifest_response = requests.get(presigned_download_url, timeout=60)
manifest_response.raise_for_status()
manifest = manifest_response.json()
# Download manifest file first
manifest_bytes = brain_client.get_file_from_storage(storage_id, filename)
manifest = json.loads(manifest_bytes.decode('utf-8'))
logging.info(
f"ℹ️\u2009 Nilor-Nodes: Processing manifest for '{manifest.get('input_name')}' with {len(manifest.get('files', []))} assets."
@@ -87,14 +96,19 @@ class MediaStreamInput:
manifest.get("files", []), key=lambda x: x.get("sequence", 0)
)
# Download all assets in parallel
# Download all assets using Brain API client
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:
# Each file_info should now contain storage_id and filename instead of presigned_url
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}")
file_bytes = brain_client.get_file_from_storage(file_storage_id, file_filename)
asset_responses.append(file_bytes)
except Exception as e:
logging.error(
f"🛑\u2009 Nilor-Nodes: Failed to download asset {file_info.get('filename')}: {e}"
)
@@ -103,9 +117,7 @@ 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
media_bytes = brain_client.get_file_from_storage(storage_id, filename)
if format == "video":
return self._process_video(media_bytes)
@@ -117,14 +129,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 +226,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 +237,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 +250,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 +261,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 or not storage_result.get('storage_id'):
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 instead of object key
storage_id = storage_result['storage_id']
final_outputs_for_sqs = {output_name: storage_id}
# After upload, send the filtered dictionary of outputs to the SQS queue.
completion_message = {
@@ -330,9 +319,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 +332,10 @@ 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"
return brain_client.upload_fileobj_to_storage(buffer, filename, "image/png")
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 +349,9 @@ 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"
return brain_client.upload_fileobj_to_storage(buffer, filename, "video/mp4")
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 +361,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)",
}