nodes/lakkiy-ComfyUI-RWImageS3
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit: ComfyUI-RWImageS3 v3.0.0

Browse Source
This commit is contained in:
lakkiy
2026-03-05 19:49:51 +08:00
commit 9ceafbd13d
7 changed files with 741 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env
+39
View File
@@ -0,0 +1,39 @@
# AWS credentials
AWS_ACCESS_KEY_ID= # Your AWS access key ID
AWS_SECRET_ACCESS_KEY= # Your AWS secret access key
# AWS region where your resources (e.g., S3 bucket) are located
AWS_REGION= # e.g., us-east-1
# S3 bucket configuration
S3_BUCKET_NAME= # Name of your S3 bucket
# Custom S3-compatible endpoint (optional, e.g., Cloudflare R2)
# ENDPOINT_URL= # e.g., https://your-account-id.r2.cloudflarestorage.com
# Fallback S3 configuration (optional)
# If image is not found in primary storage, will try fallback storage
FALLBACK_AWS_ACCESS_KEY_ID= # Fallback AWS access key ID
FALLBACK_AWS_SECRET_ACCESS_KEY= # Fallback AWS secret access key
FALLBACK_AWS_REGION= # Fallback AWS region (e.g., us-east-1)
FALLBACK_S3_BUCKET_NAME= # Fallback S3 bucket name
FALLBACK_ENDPOINT_URL= # Fallback endpoint URL (optional)
############################################################
# Example:
#
# Primary (Cloudflare R2):
# AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE
# AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY
# AWS_REGION=auto
# S3_BUCKET_NAME=my-r2-bucket
# ENDPOINT_URL=https://your-account-id.r2.cloudflarestorage.com
#
# Fallback (AWS S3):
# FALLBACK_AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE2
# FALLBACK_AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY2
# FALLBACK_AWS_REGION=us-east-1
# FALLBACK_S3_BUCKET_NAME=my-s3-bucket
# (No FALLBACK_ENDPOINT_URL for standard AWS S3)
############################################################
CLAUDE.md
+91
View File
@@ -0,0 +1,91 @@
# CLAUDE.md
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
## Project Overview
This is a ComfyUI custom node extension that provides AWS S3 integration for image processing workflows. The project enables users to read images directly from S3 storage and save processed images back to S3, eliminating the need for manual file transfers.
## Dependencies Installation
The project requires the following Python dependencies:
```bash
pip install boto3 torch pillow python-dotenv numpy
```
## Configuration
### Environment Setup
1. Create a `.env` file in the root directory with:
```env
AWS_ACCESS_KEY_ID=your_access_key_here
AWS_SECRET_ACCESS_KEY=your_secret_key_here
AWS_REGION=us-east-1
S3_BUCKET_NAME=your-bucket-name
```
2. Set up S3 bucket structure (flexible - users can now specify custom paths):
```
your-bucket-name/
├── any-folder/ # Images can be in any folder structure
│ └── subfolder/
└── output/ # Default output folder (still used by SaveImageToS3)
```
Note: With the updated ReadImageFromS3 node, users can specify any S3 object key directly in the ComfyUI interface, so the folder structure is now flexible.
## Code Architecture
### Main Components
- **`__init__.py`**: Simple module initialization that exports NODE_CLASS_MAPPINGS
- **`nodes.py`**: Core implementation containing two main node classes:
- `ReadImageFromS3`: Loads images from S3 using custom object keys (user-specified paths)
- `SaveImageToS3`: Saves processed image tensors to S3 output folder
### Key Architecture Details
1. **Image Processing Pipeline**:
- Images are downloaded to temporary files for PIL processing
- EXIF orientation correction is automatically applied
- All images are converted to RGB format (alpha channels discarded)
- Tensors are normalized to [0, 1] range for ComfyUI compatibility
2. **S3 Integration**:
- Uses boto3 client initialized with environment variables
- Supports standard image formats: .png, .jpg, .jpeg, .bmp, .webp, .tiff
- Custom S3 object key input (users can specify any path within the bucket)
- S3 object existence validation before processing
- Automatic timestamp generation for saved files
3. **ComfyUI Node Structure**:
- Both nodes follow ComfyUI conventions with INPUT_TYPES, RETURN_TYPES, etc.
- Category: "image"
- Input validation and cache invalidation methods implemented
- Error handling for S3 operations and image processing
### AWS IAM Permissions Required
The AWS credentials need these S3 permissions:
- `s3:ListBucket` - List available images
- `s3:GetObject` - Download images from input folder
- `s3:PutObject` - Upload processed images to output folder
## Installation in ComfyUI
1. Clone repository to ComfyUI's custom_nodes directory:
```bash
cd /ComfyUI/custom_nodes
git clone https://github.com/lakkiy/ComfyUI-RWImageS3.git
```
2. Install dependencies and restart ComfyUI
3. Nodes will appear in the "image" category
## Development Notes
- No build, test, or lint commands are configured in this project
- The codebase is purely Python with no additional tooling setup
- Temporary files are properly cleaned up after processing
- Error handling includes network failures and image processing errors
- The code follows defensive programming practices with proper exception handling
README.md
+140
View File
@@ -0,0 +1,140 @@
# ComfyUI-RWImageS3
ComfyUI custom nodes for reading images from S3 and saving images/videos to S3.
## Features
- Read an image directly from S3 using an explicit `s3_key`
- Save an IMAGE tensor to S3 using an explicit `s3_key`
- Save a local video file to S3 using an explicit `video_path` and `s3_key`
- Optional fallback S3 storage for image reads
- Optional custom S3 endpoint support (R2, MinIO, etc.)
- Upload/download timing logs in ComfyUI console
## Installation
1. Clone into ComfyUI custom nodes:
```bash
cd /ComfyUI/custom_nodes
git clone https://github.com/lakkiy/ComfyUI-RWImageS3.git
cd ComfyUI-RWImageS3
```
2. Install dependencies:
```bash
pip install -r requirements.txt
```
3. Configure `.env`:
```env
AWS_ACCESS_KEY_ID=your_access_key
AWS_SECRET_ACCESS_KEY=your_secret_key
AWS_REGION=us-east-1
S3_BUCKET_NAME=your-bucket
# Optional: custom S3-compatible endpoint
# ENDPOINT_URL=https://your-endpoint
# Optional fallback storage for reads
# FALLBACK_AWS_ACCESS_KEY_ID=...
# FALLBACK_AWS_SECRET_ACCESS_KEY=...
# FALLBACK_AWS_REGION=us-east-1
# FALLBACK_S3_BUCKET_NAME=...
# FALLBACK_ENDPOINT_URL=https://your-fallback-endpoint
```
4. Restart ComfyUI.
## Nodes
### Read Image From S3
Category: `image`
Input:
- `s3_key` (STRING): Full object key in bucket, e.g. `images/cat.png`
Output:
- `IMAGE`
Behavior:
- Downloads image from primary storage; if configured and primary fails, tries fallback storage
- Applies EXIF transpose
- Converts image to RGB tensor in `[0, 1]`
### Save Image To S3
Category: `image`
Input:
- `image` (IMAGE)
- `s3_key` (STRING): Full target object key, e.g. `results/run1/frame_0001.png`
Output:
- No outputs (`OUTPUT_NODE = True`)
Behavior:
- Converts tensor to PNG and uploads to S3
- Prints upload completion log with duration
### Save Video To S3
Category: `image`
Input:
- `video_path` (STRING, force input): Absolute local path, e.g. `/root/autodl-tmp/ComfyUI/output/clip.mp4`
- `s3_key` (STRING): Full target object key, e.g. `results/run1/clip.mp4`
Output:
- No outputs (`OUTPUT_NODE = True`)
Behavior:
- Validates absolute path and file existence
- Validates extension: `.mp4 .mov .webm .mkv .avi .gif`
- Uploads file to S3 and prints upload duration
### Is Mask Empty
Category: `image`
Input:
- `mask` (MASK)
Output:
- `BOOLEAN`
## Supported image formats (read)
- `.png`
- `.jpg`
- `.jpeg`
- `.bmp`
- `.webp`
- `.tiff`
## Required IAM permissions
- `s3:GetObject`
- `s3:PutObject`
- `s3:ListBucket` (recommended)
## Troubleshooting
- `S3_BUCKET_NAME is not configured`
- Set `S3_BUCKET_NAME` in `.env` and restart ComfyUI.
- `s3_key cannot be empty`
- Provide a full object key including folder and filename.
- `Please provide an absolute local video path`
- Pass an absolute path in `video_path`.
- `Local video file not found`
- Confirm the file exists on disk and ComfyUI has access.
## Version
Current version: `3.0.0`
__init__.py
+1
View File
@@ -0,0 +1 @@
from .nodes import NODE_CLASS_MAPPINGS
nodes.py
+465
View File
@@ -0,0 +1,465 @@
"""
ComfyUI S3 Media Nodes
Custom ComfyUI nodes for reading images from S3 and uploading images/videos to S3.
Included nodes:
- ReadImageFromS3
- SaveImageToS3
- SaveVideoToS3
- IsMaskEmptyNode
Requirements:
- boto3
- torch
- pillow
- python-dotenv
- numpy
GitHub: https://github.com/lakkiy/ComfyUI-RWImageS3
Version: 3.0.0
"""
import os
import tempfile
import time
from typing import Any, Dict, Tuple, Union
import boto3
import numpy as np
import torch
from dotenv import load_dotenv
from PIL import Image, ImageOps
# ═════════════════════════════════════════════════════════════════════════════
# CONFIGURATION AND INITIALIZATION
# ═════════════════════════════════════════════════════════════════════════════
# Load environment variables from .env file (must be located next to this script)
load_dotenv()
# AWS Configuration - Retrieved from environment variables
AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
AWS_REGION = os.getenv("AWS_REGION", "us-east-1") # Default to us-east-1
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
ENDPOINT_URL = os.getenv(
"ENDPOINT_URL"
) # Custom S3-compatible endpoint (optional, e.g., Cloudflare R2)
# Fallback S3 Configuration (optional) - Retrieved from environment variables
FALLBACK_AWS_ACCESS_KEY = os.getenv("FALLBACK_AWS_ACCESS_KEY_ID")
FALLBACK_AWS_SECRET_KEY = os.getenv("FALLBACK_AWS_SECRET_ACCESS_KEY")
FALLBACK_AWS_REGION = os.getenv("FALLBACK_AWS_REGION", "us-east-1")
FALLBACK_S3_BUCKET_NAME = os.getenv("FALLBACK_S3_BUCKET_NAME")
FALLBACK_ENDPOINT_URL = os.getenv("FALLBACK_ENDPOINT_URL")
# Supported image file extensions (single-frame formats only)
ALLOWED_EXTENSIONS = (".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff")
# Supported video file extensions for SaveVideoToS3
ALLOWED_VIDEO_EXTENSIONS = (".mp4", ".mov", ".webm", ".mkv", ".avi", ".gif")
# Initialize AWS S3 client with credentials
# If ENDPOINT_URL is set, use it for S3-compatible services (e.g., Cloudflare R2, MinIO)
s3_client_config = {
"aws_access_key_id": AWS_ACCESS_KEY,
"aws_secret_access_key": AWS_SECRET_KEY,
"region_name": AWS_REGION,
}
if ENDPOINT_URL:
s3_client_config["endpoint_url"] = ENDPOINT_URL
s3_client = boto3.client("s3", **s3_client_config)
# Initialize fallback S3 client if fallback credentials are provided
fallback_s3_client = None
fallback_s3_bucket = None
if FALLBACK_AWS_ACCESS_KEY and FALLBACK_AWS_SECRET_KEY and FALLBACK_S3_BUCKET_NAME:
fallback_s3_client_config = {
"aws_access_key_id": FALLBACK_AWS_ACCESS_KEY,
"aws_secret_access_key": FALLBACK_AWS_SECRET_KEY,
"region_name": FALLBACK_AWS_REGION,
}
if FALLBACK_ENDPOINT_URL:
fallback_s3_client_config["endpoint_url"] = FALLBACK_ENDPOINT_URL
fallback_s3_client = boto3.client("s3", **fallback_s3_client_config)
fallback_s3_bucket = FALLBACK_S3_BUCKET_NAME
print(f"Fallback S3 client initialized: bucket={FALLBACK_S3_BUCKET_NAME}")
# ═════════════════════════════════════════════════════════════════════════════
# COMFYUI NODE CLASSES
# ═════════════════════════════════════════════════════════════════════════════
class ReadImageFromS3:
"""
ComfyUI Node: Read Image From S3
This node provides functionality to load images directly from AWS S3 storage
into ComfyUI workflows using custom S3 object keys. Users can specify the full
S3 path to their image files, allowing for flexible folder structures.
The node processes images through PIL for proper orientation and format conversion,
and returns them as normalized PyTorch tensors ready for ComfyUI processing.
Features:
- Custom S3 path input (user can specify any S3 object key)
- Automatic EXIF orientation correction
- RGBA to RGB conversion with proper alpha handling
- Tensor normalization to [0, 1] range
- S3 object existence validation
- Error handling for network and file operations
Input:
- s3_key: Text input for full S3 object key (e.g., "my-folder/image.png")
Output:
- IMAGE: PyTorch tensor of shape [1, H, W, 3] with RGB values in [0, 1] range
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Define the input schema for the ComfyUI node interface.
Returns:
Dict containing the required input specifications for ComfyUI.
The s3_key parameter allows users to specify the full S3 object path.
"""
return {
"required": {
"s3_key": (
"STRING",
{"default": "input/example.png", "multiline": False},
),
}
}
# ComfyUI node metadata
CATEGORY = "image"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("image",)
FUNCTION = "load_image"
def load_image(self, s3_key: str) -> Tuple[torch.Tensor]:
"""
Load and process an image from S3 storage with fallback support.
This method downloads the specified image from S3 using the full S3 object key,
applies EXIF orientation correction, converts to RGBA format, extracts RGB channels,
and normalizes the pixel values to create a PyTorch tensor suitable for ComfyUI processing.
If the primary S3 storage fails to retrieve the image (e.g., 404 not found),
the method automatically attempts to download from the fallback S3 storage if configured.
Args:
s3_key (str): Full S3 object key path (e.g., "my-folder/image.png")
Returns:
Tuple[torch.Tensor]: Single-element tuple containing the processed image
tensor with shape [1, H, W, 3] and values in [0, 1]
Raises:
RuntimeError: If S3 download fails from both primary and fallback storage,
or if image processing encounters errors
"""
# Create temporary file with appropriate extension for PIL compatibility
file_extension = os.path.splitext(s3_key)[1]
with tempfile.NamedTemporaryFile(
delete=True, suffix=file_extension
) as temp_file:
# Try primary S3 storage first
download_successful = False
primary_error = None
try:
# Download image from primary S3 to temporary file
start_time = time.time()
s3_client.download_file(S3_BUCKET_NAME, s3_key, temp_file.name)
end_time = time.time()
download_duration = end_time - start_time
print(
f"S3 download completed (primary): s3_key='{s3_key}', duration={download_duration:.3f}s"
)
download_successful = True
except Exception as e:
primary_error = e
print(f"Primary S3 download failed for '{s3_key}': {e}")
# Try fallback S3 storage if configured
if fallback_s3_client is not None:
try:
print(f"Attempting fallback S3 download for '{s3_key}'...")
start_time = time.time()
fallback_s3_client.download_file(
fallback_s3_bucket, s3_key, temp_file.name
)
end_time = time.time()
download_duration = end_time - start_time
print(
f"S3 download completed (fallback): s3_key='{s3_key}', duration={download_duration:.3f}s"
)
download_successful = True
except Exception as fallback_e:
print(
f"Fallback S3 download also failed for '{s3_key}': {fallback_e}"
)
raise RuntimeError(
f"ReadImageFromS3: Failed to download '{s3_key}' from both primary and fallback storage. "
f"Primary error: {primary_error}. Fallback error: {fallback_e}"
)
else:
# No fallback configured, raise the primary error
raise RuntimeError(
f"ReadImageFromS3: Failed to download '{s3_key}': {primary_error}"
)
if not download_successful:
raise RuntimeError(
f"ReadImageFromS3: Failed to download '{s3_key}': {primary_error}"
)
# Load and process the image using PIL
try:
pil_image = Image.open(temp_file.name)
# Apply EXIF orientation correction to handle rotated images
pil_image = ImageOps.exif_transpose(pil_image)
# Convert to RGBA to ensure consistent 4-channel format
rgba_image = pil_image.convert("RGBA")
# Convert PIL image to numpy array and normalize to [0, 1]
image_array = (
np.array(rgba_image).astype(np.float32) / 255.0
) # Shape: [H, W, 4]
# Extract RGB channels (ignore alpha channel)
rgb_array = image_array[:, :, :3] # Shape: [H, W, 3]
# Convert to PyTorch tensor and add batch dimension
image_tensor = torch.from_numpy(rgb_array)[
None, ...
] # Shape: [1, H, W, 3]
return (image_tensor,)
except Exception as e:
raise RuntimeError(
f"ReadImageFromS3: Failed to process image '{s3_key}': {e}"
)
@classmethod
def IS_CHANGED(cls, s3_key: str) -> str:
"""
ComfyUI cache invalidation method.
This method helps ComfyUI determine when to re-execute the node.
Returning the S3 key ensures the node re-executes when
a different S3 object is specified.
Args:
s3_key (str): The S3 object key
Returns:
str: The S3 key for cache comparison
"""
return s3_key
@classmethod
def VALIDATE_INPUTS(cls, s3_key: str) -> Union[bool, str]:
"""
Validate that the S3 key is not empty.
This method only performs basic validation to allow lazy evaluation
scenarios where the node may not be executed (e.g., switch nodes).
S3 existence checks are performed at execution time in load_image.
Args:
s3_key (str): The S3 object key to validate
Returns:
bool: True if validation passes
str: Error message if validation fails
"""
if not s3_key or not s3_key.strip():
return "S3 key cannot be empty."
return True
class SaveImageToS3:
"""Save an IMAGE tensor to S3 using an explicit s3_key (folder + filename)."""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"image": ("IMAGE",),
"s3_key": (
"STRING",
{"default": "images/result.png", "multiline": False},
),
}
}
CATEGORY = "image"
RETURN_TYPES = ()
RETURN_NAMES = ()
FUNCTION = "save_to_s3"
OUTPUT_NODE = True
def save_to_s3(self, image: torch.Tensor, s3_key: str):
if not isinstance(image, torch.Tensor):
raise ValueError("SaveImageToS3: 'image' parameter must be a torch.Tensor")
s3_key = (s3_key or "").strip().lstrip("/")
if not s3_key:
raise RuntimeError("SaveImageToS3: s3_key cannot be empty.")
if not S3_BUCKET_NAME or not S3_BUCKET_NAME.strip():
raise RuntimeError(
"SaveImageToS3: S3_BUCKET_NAME is not configured. Please set it in your .env file."
)
processed_tensor = image[0] if image.ndim == 4 else image
image_array = processed_tensor.detach().cpu().numpy()
image_array = np.clip(image_array * 255.0, 0, 255).astype(np.uint8)
pil_image = Image.fromarray(image_array)
temp_file_path = None
try:
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
pil_image.save(temp_file.name, format="PNG")
temp_file_path = temp_file.name
start_time = time.time()
s3_client.upload_file(temp_file_path, S3_BUCKET_NAME, s3_key)
end_time = time.time()
upload_duration = end_time - start_time
print(
f"S3 upload completed (image): s3_key='{s3_key}', duration={upload_duration:.3f}s"
)
except Exception as e:
raise RuntimeError(f"SaveImageToS3: Failed to upload to '{s3_key}': {e}")
finally:
if temp_file_path and os.path.exists(temp_file_path):
os.remove(temp_file_path)
return ()
class SaveVideoToS3:
"""Upload a local video file to S3 using an explicit s3_key (folder + filename)."""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {
"required": {
"video_path": (
"STRING",
{
"forceInput": True,
},
),
"s3_key": (
"STRING",
{"default": "videos/result.mp4", "multiline": False},
),
}
}
CATEGORY = "image"
RETURN_TYPES = ()
RETURN_NAMES = ()
FUNCTION = "save_video_to_s3"
OUTPUT_NODE = True
def save_video_to_s3(self, video_path: str, s3_key: str):
if not S3_BUCKET_NAME or not S3_BUCKET_NAME.strip():
raise RuntimeError(
"SaveVideoToS3: S3_BUCKET_NAME is not configured. Please set it in your .env file."
)
local_video_path = os.path.expanduser(str(video_path).strip().strip("\"'"))
s3_key = (s3_key or "").strip().lstrip("/")
if not s3_key:
raise RuntimeError("SaveVideoToS3: s3_key cannot be empty.")
if not os.path.isabs(local_video_path):
raise RuntimeError(
"SaveVideoToS3: Please provide an absolute local video path, "
"e.g. /root/autodl-tmp/ComfyUI/output/your_video.mp4"
)
if not os.path.isfile(local_video_path):
raise RuntimeError(
f"SaveVideoToS3: Local video file not found: '{local_video_path}'"
)
ext = os.path.splitext(local_video_path)[1].lower()
if ext and ext not in ALLOWED_VIDEO_EXTENSIONS:
raise RuntimeError(
f"SaveVideoToS3: Unsupported video extension '{ext}'. "
f"Allowed: {ALLOWED_VIDEO_EXTENSIONS}"
)
try:
start_time = time.time()
s3_client.upload_file(local_video_path, S3_BUCKET_NAME, s3_key)
end_time = time.time()
upload_duration = end_time - start_time
print(
f"S3 upload completed (video): s3_key='{s3_key}', duration={upload_duration:.3f}s"
)
except Exception as e:
raise RuntimeError(
f"SaveVideoToS3: Failed to upload '{local_video_path}' to '{s3_key}': {e}"
)
return ()
# ═════════════════════════════════════════════════════════════════════════════
# UTILITY NODES
# ═════════════════════════════════════════════════════════════════════════════
class IsMaskEmptyNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
},
}
RETURN_TYPES = ("BOOLEAN",)
RETURN_NAMES = ("is_empty",)
FUNCTION = "main"
def main(self, mask):
return (bool(torch.all(mask == 0)),)
# ═════════════════════════════════════════════════════════════════════════════
# COMFYUI NODE REGISTRATION
# ═════════════════════════════════════════════════════════════════════════════
# Node class mappings for ComfyUI discovery
# This dictionary maps the display names shown in ComfyUI to the actual node classes
NODE_CLASS_MAPPINGS = {
"Read Image From S3": ReadImageFromS3,
"Save Image To S3": SaveImageToS3,
"Save Video To S3": SaveVideoToS3,
"Is Mask Empty": IsMaskEmptyNode,
}
# Optional: Display name mappings (if you want different names in the UI)
# NODE_DISPLAY_NAME_MAPPINGS = {
# "ReadImageFromS3": "Read Image From S3",
# "SaveImageToS3": "Save Image To S3",
# }
requirements.txt
+5
View File
@@ -0,0 +1,5 @@
boto3
numpy
pillow
python-dotenv
torch
workflow-example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB