Compare commits
14
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
010cd6c33b | ||
|
|
a29f1d411c | ||
|
|
cfd8e2cfde | ||
|
|
a491e6463c | ||
|
|
5b8184cc33 | ||
|
|
6b2b482b89 | ||
|
|
ea13b61f21 | ||
|
|
663d2dd739 | ||
|
|
b940ad79ea | ||
|
|
baba20c3c4 | ||
|
|
f3dbe8a3ed | ||
|
|
cb9a98edf1 | ||
|
|
e8ee057dbe | ||
|
|
5f7d00560b |
@@ -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/
|
||||
|
||||
@@ -2,6 +2,10 @@
|
||||
|
||||
A collection of utility nodes for ComfyUI focusing on list manipulation, batch operations, and advanced I/O functionality.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- `comfyui-kjnodes` custom_nodes repo
|
||||
|
||||
## 🏭 Generators
|
||||
|
||||
<details>
|
||||
|
||||
@@ -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
|
||||
+131
-77
@@ -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",
|
||||
@@ -251,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,
|
||||
@@ -264,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 = {
|
||||
@@ -331,7 +353,7 @@ class MediaStreamOutput:
|
||||
|
||||
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..."
|
||||
)
|
||||
@@ -342,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)}"
|
||||
)
|
||||
@@ -358,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 ---
|
||||
@@ -390,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)",
|
||||
}
|
||||
|
||||
+212
-45
@@ -21,6 +21,19 @@ from .utils import pil2tensor, tensor2pil
|
||||
import logging
|
||||
from comfy.utils import common_upscale
|
||||
from comfy import model_management
|
||||
import sys
|
||||
from os.path import dirname, join
|
||||
|
||||
# Attempt to import ImagePadKJ from comfyui-kjnodes if available
|
||||
_kj_nodes_path = join(dirname(__file__), "..", "comfyui-kjnodes", "nodes")
|
||||
if _kj_nodes_path not in sys.path:
|
||||
sys.path.append(_kj_nodes_path)
|
||||
try:
|
||||
from image_nodes import ImagePadKJ # type: ignore
|
||||
except Exception as _e:
|
||||
logging.warning(
|
||||
f"⚠️\u2009 Nilor-Nodes (nilornodes): Could not import ImagePadKJ from comfyui-kjnodes ({_kj_nodes_path}): {_e}"
|
||||
)
|
||||
|
||||
BIGMIN = -(2**53 - 1)
|
||||
BIGMAX = 2**53 - 1
|
||||
@@ -171,7 +184,7 @@ class NilorRemapFloatList:
|
||||
# Avoid division by zero
|
||||
if max_input - min_input == 0:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (RemapFloatList): max_input and min_input cannot be the same value."
|
||||
"🛑\u2009 Nilor-Nodes (RemapFloatList): max_input and min_input cannot be the same value."
|
||||
)
|
||||
|
||||
scale = (max_output - min_output) / (max_input - min_input)
|
||||
@@ -228,7 +241,7 @@ class NilorInverseMapFloatList:
|
||||
def inverse_map_float_list(self, list_of_floats):
|
||||
if not list_of_floats:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (InverseMapFloatList): The input list_of_floats cannot be empty."
|
||||
"🛑\u2009 Nilor-Nodes (InverseMapFloatList): The input list_of_floats cannot be empty."
|
||||
)
|
||||
|
||||
min_input = min(list_of_floats)
|
||||
@@ -326,7 +339,7 @@ class NilorCountImagesInDirectory:
|
||||
def count_images_in_directory(self, directory):
|
||||
if not os.path.isdir(directory):
|
||||
raise FileNotFoundError(
|
||||
f"[🛑] Nilor-Nodes (NilorCountImagesInDirectory): Directory '{directory} cannot be found."
|
||||
f"🛑\u2009 Nilor-Nodes (NilorCountImagesInDirectory): Directory '{directory}' cannot be found."
|
||||
)
|
||||
|
||||
list_dir = []
|
||||
@@ -376,7 +389,7 @@ class NilorSelectIndexFromList:
|
||||
# Ensure the index is within bounds
|
||||
if index < 0 or index >= len(actual_list):
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (SelectIndexFromList): Index is outside the bounds of the array."
|
||||
"🛑\u2009 Nilor-Nodes (SelectIndexFromList): Index is outside the bounds of the array."
|
||||
)
|
||||
|
||||
# Returns the value at the given index
|
||||
@@ -469,7 +482,7 @@ class NilorSaveEXRArbitrary:
|
||||
for tensor in image_channels:
|
||||
if tensor.shape[-2:] != (height, width):
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (SaveEXRArbitrary): All input tensors must have the same dimensions"
|
||||
"🛑\u2009 Nilor-Nodes (SaveEXRArbitrary): All input tensors must have the same dimensions"
|
||||
)
|
||||
|
||||
# Channel naming
|
||||
@@ -521,7 +534,7 @@ class NilorSaveEXRArbitrary:
|
||||
exr_file.close()
|
||||
|
||||
logging.info(
|
||||
f"✅ Nilor-Nodes (SaveEXRArbitrary): EXR file saved successfully to {writepath}"
|
||||
f"✅\u2009 Nilor-Nodes (SaveEXRArbitrary): EXR file saved successfully to {writepath}"
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
@@ -645,13 +658,13 @@ class NilorShuffleImageBatch:
|
||||
def _check_image_dimensions(self, images):
|
||||
if images.shape[0] == 0:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (ShuffleImageBatch): Input images tensor is empty."
|
||||
"🛑\u2009 Nilor-Nodes (ShuffleImageBatch): Input images tensor is empty."
|
||||
)
|
||||
|
||||
# All images in the batch should have the same dimensions
|
||||
if len(images.shape) != 4:
|
||||
raise ValueError(
|
||||
f"[🛑] Nilor-Nodes (ShuffleImageBatch): Expected 4D tensor (batch, channels, height, width), got shape {images.shape}"
|
||||
f"🛑\u2009 Nilor-Nodes (ShuffleImageBatch): Expected 4D tensor (batch, channels, height, width), got shape {images.shape}"
|
||||
)
|
||||
|
||||
def shuffle_image_batch(self, images: torch.Tensor, seed):
|
||||
@@ -692,13 +705,13 @@ class NilorRepeatTrimImageBatch:
|
||||
def _check_image_dimensions(self, images):
|
||||
if images.shape[0] == 0:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (RepeatTrimImageBatch): Input images tensor is empty."
|
||||
"🛑\u2009 Nilor-Nodes (RepeatTrimImageBatch): Input images tensor is empty."
|
||||
)
|
||||
|
||||
# All images in the batch should have the same dimensions
|
||||
if len(images.shape) != 4:
|
||||
raise ValueError(
|
||||
f"[🛑] Nilor-Nodes (RepeatTrimImageBatch): Expected 4D tensor (batch, channels, height, width), got shape {images.shape}"
|
||||
f"🛑\u2009 Nilor-Nodes (RepeatTrimImageBatch): Expected 4D tensor (batch, channels, height, width), got shape {images.shape}"
|
||||
)
|
||||
|
||||
def repeat_trim_image_batch(self, images: torch.Tensor, count):
|
||||
@@ -737,13 +750,13 @@ class NilorRepeatShuffleTrimImageBatch:
|
||||
def _check_image_dimensions(self, images):
|
||||
if images.shape[0] == 0:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (RepeatShuffleTrimImageBatch): Input images tensor is empty."
|
||||
"🛑\u2009 Nilor-Nodes (RepeatShuffleTrimImageBatch): Input images tensor is empty."
|
||||
)
|
||||
|
||||
# All images in the batch should have the same dimensions
|
||||
if len(images.shape) != 4:
|
||||
raise ValueError(
|
||||
f"[🛑] Nilor-Nodes (RepeatShuffleTrimImageBatch): Expected 4D tensor (batch, channels, height, width), got shape {images.shape}"
|
||||
f"🛑\u2009 Nilor-Nodes (RepeatShuffleTrimImageBatch): Expected 4D tensor (batch, channels, height, width), got shape {images.shape}"
|
||||
)
|
||||
|
||||
def repeat_shuffle_trim_image_batch(self, images: torch.Tensor, seed, count):
|
||||
@@ -870,7 +883,7 @@ class NilorNFractionsOfInt:
|
||||
return ([i * numerator // (denominator - 1) for i in range(denominator)],)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"[🛑] Nilor-Nodes (NilorNFractionsOfInt): Unknown type: {type}"
|
||||
f"🛑\u2009 Nilor-Nodes (NilorNFractionsOfInt): Unknown type: {type}"
|
||||
)
|
||||
|
||||
|
||||
@@ -951,17 +964,17 @@ class NilorWanTileResolution:
|
||||
for name, value in dims.items():
|
||||
if value <= 0:
|
||||
raise ValueError(
|
||||
f"[🛑] Nilor-Nodes (NilorWanTileResolution): {name} must be a positive integer."
|
||||
f"🛑\u2009 Nilor-Nodes (NilorWanTileResolution): {name} must be a positive integer."
|
||||
)
|
||||
|
||||
if input_width % 16 != 0 or input_height % 16 != 0:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (NilorWanTileResolution): input_width and input_height must be multiples of 16."
|
||||
"🛑\u2009 Nilor-Nodes (NilorWanTileResolution): input_width and input_height must be multiples of 16."
|
||||
)
|
||||
|
||||
if target_width < self.MIN_TILE_DIM or target_height < self.MIN_TILE_DIM:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (NilorWanTileResolution): target_width and target_height must be at least the minimum tile size."
|
||||
"🛑\u2009 Nilor-Nodes (NilorWanTileResolution): target_width and target_height must be at least the minimum tile size."
|
||||
)
|
||||
|
||||
min_blocks = self.MIN_TILE_DIM // 16
|
||||
@@ -972,7 +985,7 @@ class NilorWanTileResolution:
|
||||
|
||||
if max_width_blocks < min_blocks or max_height_blocks < min_blocks:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (NilorWanTileResolution): Target dimensions do not allow a tile within the supported range."
|
||||
"🛑\u2009 Nilor-Nodes (NilorWanTileResolution): Target dimensions do not allow a tile within the supported range."
|
||||
)
|
||||
|
||||
aspect_ratio = input_width / input_height
|
||||
@@ -1018,12 +1031,59 @@ class NilorWanTileResolution:
|
||||
if best_dimensions is None:
|
||||
# If no suitable tile resolution was found, raise an error
|
||||
raise RuntimeError(
|
||||
"[🛑] Nilor-Nodes (NilorWanTileResolution): Failed to determine a suitable tile resolution."
|
||||
"🛑\u2009 Nilor-Nodes (NilorWanTileResolution): Failed to determine a suitable tile resolution."
|
||||
)
|
||||
|
||||
return best_dimensions
|
||||
|
||||
|
||||
class NilorWanFrameTrim:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("images",)
|
||||
FUNCTION = "trim_to_wan_count"
|
||||
CATEGORY = category + subcategories["utilities"]
|
||||
|
||||
def _validate_images(self, images):
|
||||
if not isinstance(images, torch.Tensor):
|
||||
raise TypeError(
|
||||
"🛑\u2009 Nilor-Nodes (WanFrameTrim): images must be a torch.Tensor."
|
||||
)
|
||||
if images.dim() != 4:
|
||||
raise ValueError(
|
||||
f"🛑\u2009 Nilor-Nodes (WanFrameTrim): Expected 4D tensor (batch, height, width, channels), got shape {tuple(images.shape)}"
|
||||
)
|
||||
if images.shape[0] == 0:
|
||||
raise ValueError(
|
||||
"🛑\u2009 Nilor-Nodes (WanFrameTrim): Input images tensor is empty."
|
||||
)
|
||||
|
||||
def trim_to_wan_count(self, images: torch.Tensor):
|
||||
self._validate_images(images)
|
||||
|
||||
batch_count = images.shape[0]
|
||||
# Find the largest m <= batch_count such that m ≡ 1 (mod 4)
|
||||
wan_count = batch_count - ((batch_count - 1) % 4)
|
||||
|
||||
if wan_count <= 0:
|
||||
raise ValueError(
|
||||
"🛑\u2009 Nilor-Nodes (WanFrameTrim): Unable to compute a valid 4N+1 frame count from input."
|
||||
)
|
||||
|
||||
trimmed = images[:wan_count]
|
||||
return (trimmed,)
|
||||
|
||||
|
||||
class NilorCategorizeString:
|
||||
def __init__(self):
|
||||
pass
|
||||
@@ -1144,7 +1204,7 @@ class NilorRandomString:
|
||||
]
|
||||
if not options:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (NilorRandomString): No valid choices provided."
|
||||
"🛑\u2009 Nilor-Nodes (NilorRandomString): No valid choices provided."
|
||||
)
|
||||
|
||||
# Limit to the first 'max_options' entries if there are more options
|
||||
@@ -1188,7 +1248,7 @@ class NilorLoadImageByIndex:
|
||||
def load_image_by_index(self, image_directory, seed, sort_mode, reverse_sort):
|
||||
if not os.path.exists(image_directory):
|
||||
raise FileNotFoundError(
|
||||
f"[🛑] Nilor-Nodes (NilorLoadImageByIndex): Image directory {image_directory} does not exist"
|
||||
f"🛑\u2009 Nilor-Nodes (NilorLoadImageByIndex): Image directory {image_directory} does not exist"
|
||||
)
|
||||
|
||||
# Get list of image files
|
||||
@@ -1202,7 +1262,7 @@ class NilorLoadImageByIndex:
|
||||
|
||||
if not files:
|
||||
raise ValueError(
|
||||
f"[🛑] Nilor-Nodes (NilorLoadImageByIndex): No image files found in {image_directory}"
|
||||
f"🛑\u2009 Nilor-Nodes (NilorLoadImageByIndex): No image files found in {image_directory}"
|
||||
)
|
||||
|
||||
# Sort files based on selected mode
|
||||
@@ -1254,7 +1314,7 @@ class NilorExtractFilenameFromPath:
|
||||
# Ensure the input is a valid path
|
||||
if not filepath:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (ExtractFilenameFromPath): Filepath cannot be empty."
|
||||
"🛑\u2009 Nilor-Nodes (ExtractFilenameFromPath): Filepath cannot be empty."
|
||||
)
|
||||
|
||||
path = Path(filepath)
|
||||
@@ -1291,7 +1351,7 @@ class NilorBlurAnalysis:
|
||||
# Ensure images is a 4D tensor.
|
||||
if images.dim() != 4:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (BlurAnalysis): Input images must be a 4D tensor (batch, channels/height, height/width, width/channels)"
|
||||
"🛑\u2009 Nilor-Nodes (BlurAnalysis): Input images must be a 4D tensor (batch, channels/height, height/width, width/channels)"
|
||||
)
|
||||
|
||||
# Detect if using NCHW or NHWC.
|
||||
@@ -1300,7 +1360,7 @@ class NilorBlurAnalysis:
|
||||
images = images.permute(0, 3, 1, 2)
|
||||
else:
|
||||
raise ValueError(
|
||||
"[🛑] Nilor-Nodes (BlurAnalysis): Cannot determine image format (expected channel to be 1 or 3)."
|
||||
"🛑\u2009 Nilor-Nodes (BlurAnalysis): Cannot determine image format (expected channel to be 1 or 3)."
|
||||
)
|
||||
|
||||
output_images = []
|
||||
@@ -1382,6 +1442,7 @@ class NilorToSparseIndexMethod:
|
||||
|
||||
class NilorImageResizeV2:
|
||||
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
@@ -1390,15 +1451,41 @@ class NilorImageResizeV2:
|
||||
"width": ("INT", {"default": 512, "min": 0, "max": BIGMAX, "step": 1}),
|
||||
"height": ("INT", {"default": 512, "min": 0, "max": BIGMAX, "step": 1}),
|
||||
"upscale_method": (s.upscale_methods,),
|
||||
"keep_proportion": (["stretch", "resize", "pad", "pad_edge", "pad_edge_pixel", "crop", "pillarbox_blur"], {"default": False}),
|
||||
"keep_proportion": (
|
||||
[
|
||||
"stretch",
|
||||
"resize",
|
||||
"pad",
|
||||
"pad_edge",
|
||||
"pad_edge_pixel",
|
||||
"crop",
|
||||
"pillarbox_blur",
|
||||
],
|
||||
{"default": False},
|
||||
),
|
||||
"pad_color": ("STRING", {"default": "0, 0, 0"}),
|
||||
"crop_position": (["center", "top", "bottom", "left", "right"], {"default": "center"}),
|
||||
"divisible_by": ("INT", {"default": 2, "min": 0, "max": 512, "step": 1}),
|
||||
"crop_position": (
|
||||
["center", "top", "bottom", "left", "right"],
|
||||
{"default": "center"},
|
||||
),
|
||||
"divisible_by": (
|
||||
"INT",
|
||||
{"default": 2, "min": 0, "max": 512, "step": 1},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK",),
|
||||
"device": (["cpu", "gpu"],),
|
||||
"per_batch": ("INT", {"default": 16, "min": 0, "max": 4096, "step": 1, "tooltip": "Process images in sub-batches. 0 disables."}),
|
||||
"per_batch": (
|
||||
"INT",
|
||||
{
|
||||
"default": 16,
|
||||
"min": 0,
|
||||
"max": 4096,
|
||||
"step": 1,
|
||||
"tooltip": "Process images in sub-batches. 0 disables.",
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {"unique_id": "UNIQUE_ID"},
|
||||
}
|
||||
@@ -1411,12 +1498,28 @@ class NilorImageResizeV2:
|
||||
Resizes images with optional aspect preservation, padding/cropping, and sub-batching to lower peak memory.
|
||||
"""
|
||||
|
||||
def resize(self, image, width, height, keep_proportion, upscale_method, divisible_by, pad_color, crop_position, unique_id, device="cpu", mask=None, per_batch=16):
|
||||
def resize(
|
||||
self,
|
||||
image,
|
||||
width,
|
||||
height,
|
||||
keep_proportion,
|
||||
upscale_method,
|
||||
divisible_by,
|
||||
pad_color,
|
||||
crop_position,
|
||||
unique_id,
|
||||
device="cpu",
|
||||
mask=None,
|
||||
per_batch=16,
|
||||
):
|
||||
B, H, W, C = image.shape
|
||||
|
||||
if device == "gpu":
|
||||
if upscale_method == "lanczos":
|
||||
raise Exception("Lanczos is not supported on the GPU")
|
||||
raise Exception(
|
||||
"🛑\u2009 Nilor-Nodes (NilorImageResizeV2): Lanczos is not supported on the GPU"
|
||||
)
|
||||
device = model_management.get_torch_device()
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
@@ -1427,7 +1530,11 @@ Resizes images with optional aspect preservation, padding/cropping, and sub-batc
|
||||
height = H
|
||||
|
||||
pillarbox_blur = keep_proportion == "pillarbox_blur"
|
||||
if keep_proportion == "resize" or keep_proportion.startswith("pad") or pillarbox_blur:
|
||||
if (
|
||||
keep_proportion == "resize"
|
||||
or keep_proportion.startswith("pad")
|
||||
or pillarbox_blur
|
||||
):
|
||||
if width == 0 and height != 0:
|
||||
ratio = height / H
|
||||
new_width = round(W * ratio)
|
||||
@@ -1484,13 +1591,19 @@ Resizes images with optional aspect preservation, padding/cropping, and sub-batc
|
||||
bytes_per_elem = image.element_size()
|
||||
est_total_bytes = B * height * width * C * bytes_per_elem
|
||||
est_mb = est_total_bytes / (1024 * 1024)
|
||||
print(f"[NilorImageResizeV2] estimated output ~{est_mb:.2f} MB; batching {per_batch}/{B}")
|
||||
logging.info(
|
||||
f"ℹ️\u2009 Nilor-Nodes (NilorImageResizeV2) Estimated output ~{est_mb:.2f} MB."
|
||||
)
|
||||
except:
|
||||
pass
|
||||
|
||||
def _process_subbatch(in_image, in_mask):
|
||||
out_image = in_image if in_image.device == device else in_image.to(device)
|
||||
out_mask = None if in_mask is None else (in_mask if in_mask.device == device else in_mask.to(device))
|
||||
out_mask = (
|
||||
None
|
||||
if in_mask is None
|
||||
else (in_mask if in_mask.device == device else in_mask.to(device))
|
||||
)
|
||||
|
||||
if keep_proportion == "crop":
|
||||
old_height = out_image.shape[-3]
|
||||
@@ -1522,14 +1635,30 @@ Resizes images with optional aspect preservation, padding/cropping, and sub-batc
|
||||
if out_mask is not None:
|
||||
out_mask = out_mask.narrow(-1, x, crop_w).narrow(-2, y, crop_h)
|
||||
|
||||
out_image = common_upscale(out_image.movedim(-1, 1), width, height, upscale_method, crop="disabled").movedim(1, -1)
|
||||
out_image = common_upscale(
|
||||
out_image.movedim(-1, 1), width, height, upscale_method, crop="disabled"
|
||||
).movedim(1, -1)
|
||||
if out_mask is not None:
|
||||
if upscale_method == "lanczos":
|
||||
out_mask = common_upscale(out_mask.unsqueeze(1).repeat(1, 3, 1, 1), width, height, upscale_method, crop="disabled").movedim(1, -1)[:, :, :, 0]
|
||||
out_mask = common_upscale(
|
||||
out_mask.unsqueeze(1).repeat(1, 3, 1, 1),
|
||||
width,
|
||||
height,
|
||||
upscale_method,
|
||||
crop="disabled",
|
||||
).movedim(1, -1)[:, :, :, 0]
|
||||
else:
|
||||
out_mask = common_upscale(out_mask.unsqueeze(1), width, height, upscale_method, crop="disabled").squeeze(1)
|
||||
out_mask = common_upscale(
|
||||
out_mask.unsqueeze(1),
|
||||
width,
|
||||
height,
|
||||
upscale_method,
|
||||
crop="disabled",
|
||||
).squeeze(1)
|
||||
|
||||
if (keep_proportion.startswith("pad") or pillarbox_blur) and (pad_left > 0 or pad_right > 0 or pad_top > 0 or pad_bottom > 0):
|
||||
if (keep_proportion.startswith("pad") or pillarbox_blur) and (
|
||||
pad_left > 0 or pad_right > 0 or pad_top > 0 or pad_bottom > 0
|
||||
):
|
||||
padded_width = width + pad_left + pad_right
|
||||
padded_height = height + pad_top + pad_bottom
|
||||
if divisible_by > 1:
|
||||
@@ -1543,12 +1672,30 @@ Resizes images with optional aspect preservation, padding/cropping, and sub-batc
|
||||
pad_bottom += extra_height
|
||||
|
||||
pad_mode = (
|
||||
"pillarbox_blur" if pillarbox_blur else
|
||||
"edge" if keep_proportion == "pad_edge" else
|
||||
"edge_pixel" if keep_proportion == "pad_edge_pixel" else
|
||||
"color"
|
||||
"pillarbox_blur"
|
||||
if pillarbox_blur
|
||||
else (
|
||||
"edge"
|
||||
if keep_proportion == "pad_edge"
|
||||
else (
|
||||
"edge_pixel"
|
||||
if keep_proportion == "pad_edge_pixel"
|
||||
else "color"
|
||||
)
|
||||
)
|
||||
)
|
||||
out_image, out_mask = ImagePadKJ.pad(
|
||||
self,
|
||||
out_image,
|
||||
pad_left,
|
||||
pad_right,
|
||||
pad_top,
|
||||
pad_bottom,
|
||||
0,
|
||||
pad_color,
|
||||
pad_mode,
|
||||
mask=out_mask,
|
||||
)
|
||||
out_image, out_mask = ImagePadKJ.pad(self, out_image, pad_left, pad_right, pad_top, pad_bottom, 0, pad_color, pad_mode, mask=out_mask)
|
||||
|
||||
return out_image, out_mask
|
||||
|
||||
@@ -1567,9 +1714,13 @@ Resizes images with optional aspect preservation, padding/cropping, and sub-batc
|
||||
sub_out_img, sub_out_mask = _process_subbatch(sub_img, sub_mask)
|
||||
chunks.append(sub_out_img.cpu())
|
||||
if mask is not None:
|
||||
mask_chunks.append(sub_out_mask.cpu() if sub_out_mask is not None else None)
|
||||
mask_chunks.append(
|
||||
sub_out_mask.cpu() if sub_out_mask is not None else None
|
||||
)
|
||||
try:
|
||||
print(f"[NilorImageResizeV2] batch {current_batch}/{total_batches} · images {end_idx}/{B}")
|
||||
logging.info(
|
||||
f"ℹ️\u2009 Nilor-Nodes (NilorImageResizeV2) Batch {current_batch}/{total_batches} · images {end_idx}/{B}"
|
||||
)
|
||||
except:
|
||||
pass
|
||||
out_image = torch.cat(chunks, dim=0)
|
||||
@@ -1578,7 +1729,21 @@ Resizes images with optional aspect preservation, padding/cropping, and sub-batc
|
||||
else:
|
||||
out_mask = None
|
||||
|
||||
return (out_image.cpu(), out_image.shape[2], out_image.shape[1], out_mask.cpu() if out_mask is not None else torch.zeros(64, 64, device=torch.device("cpu"), dtype=torch.float32))
|
||||
logging.info(f"✅\u2009 Nilor-Nodes (NilorImageResizeV2) All batches complete.")
|
||||
|
||||
return (
|
||||
out_image.cpu(),
|
||||
out_image.shape[2],
|
||||
out_image.shape[1],
|
||||
(
|
||||
out_mask.cpu()
|
||||
if out_mask is not None
|
||||
else torch.zeros(
|
||||
64, 64, device=torch.device("cpu"), dtype=torch.float32
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# Mapping class names to objects for potential export
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
@@ -1607,6 +1772,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"Nilor Blur Analysis": NilorBlurAnalysis,
|
||||
"Nilor To Sparse Index Method": NilorToSparseIndexMethod,
|
||||
"Nilor Image Resize v2": NilorImageResizeV2,
|
||||
"Nilor Wan Frame Trim": NilorWanFrameTrim,
|
||||
}
|
||||
|
||||
# Mapping nodes to human-readable names
|
||||
@@ -1636,4 +1802,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Nilor Blur Analysis": "👺 Blur Analysis",
|
||||
"Nilor To Sparse Index Method": "👺 To Sparse Index Method",
|
||||
"Nilor Image Resize v2": "👺 Resize Image v2",
|
||||
"Nilor Wan Frame Trim": "👺 Wan Frame Trim",
|
||||
}
|
||||
|
||||
+1
-1
@@ -4,7 +4,7 @@ aiofiles>=23.2.1
|
||||
aiohttp==3.12.14
|
||||
boto3==1.40.15
|
||||
fastapi==0.110.0
|
||||
huggingface_hub==0.33.4
|
||||
huggingface_hub==0.34.0
|
||||
imageio==2.37.0
|
||||
imageio-ffmpeg==0.6.0
|
||||
numpy>=1.26.4
|
||||
|
||||
Reference in New Issue
Block a user