109 lines
4.0 KiB
Python
109 lines
4.0 KiB
Python
import os
|
|
import json
|
|
import boto3
|
|
from PIL import Image
|
|
import io
|
|
import torch
|
|
import numpy as np
|
|
import secrets
|
|
import string
|
|
import time
|
|
import random
|
|
|
|
class S3SaveNode:
|
|
def __init__(self):
|
|
self.type = "output"
|
|
self.s3_client = boto3.client(
|
|
"s3",
|
|
endpoint_url=os.environ.get("CLOUDFLARE_R2_ENDPOINT_URL", ""),
|
|
aws_access_key_id=os.environ.get("CLOUDFLARE_ACCESS_KEY_ID", ""),
|
|
aws_secret_access_key=os.environ.get("CLOUDFLARE_SECRET_ACCESS_KEY", ""),
|
|
)
|
|
self.cloudflare_r2_url = os.environ.get("CLOUDFLARE_R2_URL", "https://img-dev.fantasy.ai/")
|
|
self.cloudflare_r2_bucket = os.environ.get("CLOUDFLARE_R2_BUCKET", "fan-dev")
|
|
self.cloudflare_r2_prefix = os.environ.get("CLOUDFLARE_R2_PREFIX", "vik_tests/")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "save_to_s3"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "image/s3"
|
|
|
|
def save_to_s3(self, images):
|
|
bucket = self.cloudflare_r2_bucket
|
|
prefix = self.cloudflare_r2_prefix
|
|
results = []
|
|
|
|
try:
|
|
for i, image in enumerate(images):
|
|
try:
|
|
print(f"Processing image {i + 1}/{len(images)}")
|
|
|
|
# Convert image to numpy
|
|
if isinstance(image, torch.Tensor):
|
|
if image.ndim == 4:
|
|
image = image[0]
|
|
image = image.cpu().numpy()
|
|
if image.shape[0] == 3:
|
|
image = np.transpose(image, (1, 2, 0))
|
|
image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8)
|
|
elif isinstance(image, np.ndarray):
|
|
if image.shape[0] == 3 and image.ndim == 3:
|
|
image = np.transpose(image, (1, 2, 0))
|
|
image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8)
|
|
else:
|
|
raise ValueError("Unsupported image type")
|
|
|
|
pil_img = Image.fromarray(image)
|
|
buffer = io.BytesIO()
|
|
pil_img.save(buffer, format="WEBP", quality=95)
|
|
buffer.seek(0)
|
|
|
|
filename = ''.join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20)) + '.webp'
|
|
key = prefix + filename
|
|
|
|
print(f"Uploading to R2: {bucket}/{key}")
|
|
self.s3_client.put_object(
|
|
Bucket=bucket,
|
|
Key=key,
|
|
Body=buffer.getvalue(),
|
|
ContentType="image/webp"
|
|
)
|
|
|
|
# Build public URL using the correct public domain: CF > R2 > bucket > Settings > Custom Domains
|
|
public_url = f"{self.cloudflare_r2_url}{key}"
|
|
|
|
# Append URL and status to results
|
|
results.append({
|
|
"url": public_url,
|
|
"status": "success"
|
|
})
|
|
|
|
print(f" Uploaded: {filename}")
|
|
print(f" URL: {public_url}")
|
|
|
|
except Exception as e:
|
|
print(f" Failed to upload image {i}: {e}")
|
|
|
|
# Append failed result
|
|
results.append({
|
|
"url": None,
|
|
"status": "failed"
|
|
})
|
|
|
|
# Return JSON string of results
|
|
return {"STRING": json.dumps(results)}
|
|
|
|
except Exception as e:
|
|
error_msg = f"Fatal error: {str(e)}"
|
|
print(error_msg)
|
|
return {json.dumps([{"url": None, "status": "failed"}])}
|