Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
010cd6c33b | ||
|
|
a29f1d411c | ||
|
|
cfd8e2cfde |
@@ -1,8 +1,12 @@
|
||||
# --- BRAIN CREDENTIALS ---
|
||||
BRANDO_API_KEY=...
|
||||
|
||||
# --- S3 (AWS/MINIO) CREDENTIALS ---
|
||||
## These are used by both scripts to connect to SQS and S3 (MinIO).
|
||||
AWS_ACCESS_KEY_ID=minioadmin
|
||||
AWS_SECRET_ACCESS_KEY=...
|
||||
AWS_DEFAULT_REGION=us-east-1
|
||||
MINIO_ENDPOINT=http://127.0.0.1:9000
|
||||
|
||||
# --- SQS SETTINGS ---
|
||||
## Toggles functionality for the SQS Worker Consumer
|
||||
|
||||
+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,392 @@
|
||||
"""
|
||||
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
|
||||
+133
-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,35 @@ 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 +108,28 @@ 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 +138,16 @@ 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 +159,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 +256,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 +267,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 +280,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 +291,30 @@ 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
|
||||
|
||||
# Upload the media using Brain API client
|
||||
brain_client = get_brain_api_client()
|
||||
storage_result = None
|
||||
|
||||
# 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.
|
||||
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 +351,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 +364,36 @@ 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"
|
||||
)
|
||||
|
||||
def _upload_video(self, image_batch_tensor, url, framerate):
|
||||
# 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, brain_client, framerate, output_name):
|
||||
logging.info(
|
||||
f"ℹ️\u2009 Nilor-Nodes (MediaStreamOutput): Uploading as MP4 video. Frame count: {len(image_batch_tensor)}"
|
||||
)
|
||||
@@ -359,29 +407,34 @@ 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"
|
||||
)
|
||||
|
||||
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
|
||||
# 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
|
||||
|
||||
|
||||
# --- Node Mappings ---
|
||||
@@ -391,6 +444,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