wavespeed api node
This commit is contained in:
nofunstudio
@@ -9,6 +9,7 @@ from .ksampler_fun import KSampler
|
||||
from .fal_recraft_upscale import FalAPI_recraft_upscale
|
||||
from .fal_kling_video import FalAPI_kling_video
|
||||
from .fal_seedance_video import FalAPI_seedance_video
|
||||
from .wavespeed_ai_image import WaveSpeedAI_Image
|
||||
from .string_lower import StringLower
|
||||
from .multi_alpha_composite import MultiAlphaComposite
|
||||
from .flux_kontext_inpaint_conditioning import FluxKontextInpaintingConditioning
|
||||
@@ -26,6 +27,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"FalAPI_recraft_upscale": FalAPI_recraft_upscale,
|
||||
"FalAPI_kling_video": FalAPI_kling_video,
|
||||
"FalAPI_seedance_video": FalAPI_seedance_video,
|
||||
"WaveSpeedAI_Image": WaveSpeedAI_Image,
|
||||
"StringLower": StringLower,
|
||||
"MultiAlphaComposite": MultiAlphaComposite,
|
||||
"FluxKontextInpaintingConditioning": FluxKontextInpaintingConditioning,
|
||||
@@ -42,6 +44,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FalAPI_recraft_upscale": "Fal API Recraft Upscale",
|
||||
"FalAPI_kling_video": "Fal API Kling Video",
|
||||
"FalAPI_seedance_video": "Fal API Seedance Video",
|
||||
"WaveSpeedAI_Image": "WaveSpeedAI Image Generation",
|
||||
"StringLower": "String to Lowercase",
|
||||
"MultiAlphaComposite": "Multi Alpha Composite",
|
||||
"FluxKontextInpaintingConditioning": "No Fun Flux Kontext Inpaint Conditioning",
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
+25
-36
@@ -34,11 +34,8 @@ class LoadImageBatchFun:
|
||||
Initialize instance-specific variables to prevent shared state issues.
|
||||
This is the key fix for supporting multiple instances.
|
||||
"""
|
||||
self.current_images = []
|
||||
self.current_filenames = []
|
||||
self.current_index = 0
|
||||
self.last_hash = None
|
||||
self.cached_images = []
|
||||
self.cached_filenames = []
|
||||
|
||||
def load_batch_images(self, mode, index, label, path, pattern, allow_RGBA_output):
|
||||
@@ -58,39 +55,38 @@ class LoadImageBatchFun:
|
||||
raise ValueError(f"Path is not a directory: {path}")
|
||||
|
||||
# Generate a hash for path/pattern to determine if we need to reload images
|
||||
# Don't include index/mode in this hash so we can reuse cached images
|
||||
# Don't include index/mode in this hash so we can reuse cached image paths
|
||||
cache_hash = hashlib.md5(
|
||||
f"{path}_{pattern}_{allow_RGBA_output}".encode()
|
||||
f"{path}_{pattern}".encode()
|
||||
).hexdigest()
|
||||
|
||||
# Load images if cache is invalid or parameters changed
|
||||
if self.last_hash != cache_hash or not self.cached_images:
|
||||
print(f"[Load Image Batch Fun] Loading images from: {path}")
|
||||
self.cached_images, self.cached_filenames = self._load_images_from_path(
|
||||
path, pattern, allow_RGBA_output
|
||||
# Load image paths if cache is invalid or parameters changed
|
||||
if self.last_hash != cache_hash or not self.cached_filenames:
|
||||
print(f"[Load Image Batch Fun] Loading image paths from: {path}")
|
||||
self.cached_filenames = self._load_images_from_path(
|
||||
path, pattern
|
||||
)
|
||||
self.last_hash = cache_hash
|
||||
self.current_index = 0
|
||||
|
||||
if not self.cached_images:
|
||||
if not self.cached_filenames:
|
||||
raise ValueError(f"No images found in path: {path} with pattern: {pattern}")
|
||||
|
||||
total_images = len(self.cached_images)
|
||||
total_images = len(self.cached_filenames)
|
||||
print(f"[Load Image Batch Fun] Total images found: {total_images}")
|
||||
|
||||
# Select image based on mode
|
||||
# Select image path based on mode
|
||||
selected_path = None
|
||||
if mode == "single_image":
|
||||
# Use the specified index, wrap around if needed
|
||||
selected_index = index % total_images
|
||||
selected_image = self.cached_images[selected_index]
|
||||
selected_filename = self.cached_filenames[selected_index]
|
||||
selected_path = self.cached_filenames[selected_index]
|
||||
print(f"[Load Image Batch Fun] Single image mode - selected index: {selected_index}")
|
||||
|
||||
elif mode == "incremental_image":
|
||||
# Use current index and increment
|
||||
selected_index = self.current_index % total_images
|
||||
selected_image = self.cached_images[selected_index]
|
||||
selected_filename = self.cached_filenames[selected_index]
|
||||
selected_path = self.cached_filenames[selected_index]
|
||||
self.current_index += 1
|
||||
print(f"[Load Image Batch Fun] Incremental mode - selected index: {selected_index}, next: {self.current_index}")
|
||||
|
||||
@@ -98,19 +94,24 @@ class LoadImageBatchFun:
|
||||
# Random selection
|
||||
import random
|
||||
selected_index = random.randint(0, total_images - 1)
|
||||
selected_image = self.cached_images[selected_index]
|
||||
selected_filename = self.cached_filenames[selected_index]
|
||||
selected_path = self.cached_filenames[selected_index]
|
||||
print(f"[Load Image Batch Fun] Random mode - selected index: {selected_index}")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid mode: {mode}")
|
||||
|
||||
# Load the selected image
|
||||
selected_image = self._load_image(selected_path, allow_RGBA_output)
|
||||
|
||||
# Get filename without extension
|
||||
selected_filename = os.path.splitext(os.path.basename(selected_path))[0]
|
||||
|
||||
return (selected_image, selected_filename, total_images)
|
||||
|
||||
def _load_images_from_path(self, path, pattern, allow_RGBA_output):
|
||||
def _load_images_from_path(self, path, pattern):
|
||||
"""
|
||||
Load all images from the specified path matching the pattern.
|
||||
Returns lists of image tensors and filenames.
|
||||
Get all image file paths from the specified path matching the pattern.
|
||||
Returns a sorted list of full file paths.
|
||||
"""
|
||||
import glob
|
||||
|
||||
@@ -133,23 +134,11 @@ class LoadImageBatchFun:
|
||||
image_files.sort()
|
||||
|
||||
if not image_files:
|
||||
return [], []
|
||||
return []
|
||||
|
||||
print(f"[Load Image Batch Fun] Found {len(image_files)} image files")
|
||||
|
||||
images = []
|
||||
filenames = []
|
||||
|
||||
for img_path in image_files:
|
||||
try:
|
||||
img = self._load_image(img_path, allow_RGBA_output)
|
||||
images.append(img)
|
||||
filenames.append(os.path.basename(img_path))
|
||||
except Exception as e:
|
||||
print(f"[Load Image Batch Fun] Error loading {img_path}: {str(e)}")
|
||||
continue
|
||||
|
||||
return images, filenames
|
||||
return image_files
|
||||
|
||||
def _load_image(self, image_path, allow_RGBA_output):
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
import os
|
||||
import time
|
||||
import requests
|
||||
import tempfile
|
||||
import json
|
||||
import random
|
||||
import numpy as np
|
||||
import torch
|
||||
import asyncio
|
||||
import folder_paths
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
import base64
|
||||
|
||||
class WaveSpeedAI_Image:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
"""
|
||||
An image generation node for WavespeedAI's qwen-image/edit-plus-lora API.
|
||||
Takes up to 3 images, a prompt, and other parameters to generate a new image.
|
||||
"""
|
||||
return {
|
||||
"required": {
|
||||
"api_key": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": os.getenv("WAVESPEED_API_KEY", ""),
|
||||
"display": "WaveSpeedAI API Key"
|
||||
}),
|
||||
"prompt": ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "A beautiful painting",
|
||||
"display": "Prompt"
|
||||
}),
|
||||
"image1": ("IMAGE", {"display": "Image 1"}),
|
||||
"width": ("INT", {"default": 886, "min": 64, "max": 4096, "step": 8, "display": "Width"}),
|
||||
"height": ("INT", {"default": 1182, "min": 64, "max": 4096, "step": 8, "display": "Height"}),
|
||||
},
|
||||
"optional": {
|
||||
"image2": ("IMAGE", {"display": "Image 2 (optional)"}),
|
||||
"image3": ("IMAGE", {"display": "Image 3 (optional)"}),
|
||||
"lora_url": ("STRING", {
|
||||
"multiline": False,
|
||||
"default": "",
|
||||
"display": "LoRA URL (optional)"
|
||||
}),
|
||||
"lora_scale": ("FLOAT", {
|
||||
"default": 0.84,
|
||||
"min": 0.0,
|
||||
"max": 2.0,
|
||||
"step": 0.01,
|
||||
"display": "LoRA Scale"
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": -1,
|
||||
"min": -1,
|
||||
"max": 2147483647,
|
||||
"display": "Seed (-1 for random)"
|
||||
}),
|
||||
"output_format": (["jpeg", "png", "webp"], {
|
||||
"default": "jpeg",
|
||||
"display": "Output Format"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "STRING", "STRING",)
|
||||
RETURN_NAMES = ("image", "image_path", "generation_info",)
|
||||
FUNCTION = "generate_image"
|
||||
CATEGORY = "WaveSpeedAI"
|
||||
|
||||
def __init__(self):
|
||||
self.temp_dir = folder_paths.get_temp_directory()
|
||||
self.metadata_dir = os.path.join(self.temp_dir, "wavespeed_metadata")
|
||||
os.makedirs(self.metadata_dir, exist_ok=True)
|
||||
|
||||
def get_temp_filename(self, output_format="jpeg"):
|
||||
timestamp = int(time.time() * 1000)
|
||||
random_id = random.randint(1000, 9999)
|
||||
return f"wavespeed_{timestamp}_{random_id}.{output_format}"
|
||||
|
||||
def save_image_and_metadata(self, image_content, generation_info, filename):
|
||||
filepath = os.path.join(self.temp_dir, filename)
|
||||
with open(filepath, 'wb') as f:
|
||||
f.write(image_content)
|
||||
|
||||
metadata_filename = f"{os.path.splitext(filename)[0]}_metadata.json"
|
||||
metadata_filepath = os.path.join(self.metadata_dir, metadata_filename)
|
||||
with open(metadata_filepath, "w", encoding="utf-8") as f:
|
||||
json.dump(generation_info, f, indent=4, ensure_ascii=False)
|
||||
return filepath, metadata_filepath
|
||||
|
||||
async def generate_image(self, api_key, prompt, image1, width, height,
|
||||
image2=None, image3=None, lora_url="", lora_scale=0.84,
|
||||
seed=-1, output_format="jpeg"):
|
||||
print("[WaveSpeedAI] Starting image generation process...")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError("A WaveSpeedAI API key is required.")
|
||||
if not prompt or not prompt.strip():
|
||||
raise ValueError("A prompt is required.")
|
||||
|
||||
try:
|
||||
image_urls = []
|
||||
all_images = [image1, image2, image3]
|
||||
for i, img_tensor in enumerate(all_images):
|
||||
if img_tensor is not None:
|
||||
print(f"[WaveSpeedAI] Processing input image {i+1}...")
|
||||
data_uri = await asyncio.to_thread(self.tensor_to_data_uri, img_tensor)
|
||||
image_urls.append(data_uri)
|
||||
|
||||
if not image_urls:
|
||||
raise ValueError("At least one image input is required.")
|
||||
|
||||
payload = {
|
||||
"enable_base64_output": False,
|
||||
"enable_sync_mode": False,
|
||||
"images": image_urls,
|
||||
"output_format": output_format,
|
||||
"prompt": prompt,
|
||||
"size": f"{width}*{height}"
|
||||
}
|
||||
|
||||
if lora_url.strip():
|
||||
payload["loras"] = [{"path": lora_url, "scale": lora_scale}]
|
||||
|
||||
if seed != -1:
|
||||
payload["seed"] = seed
|
||||
|
||||
print("[WaveSpeedAI] Submitting generation task...")
|
||||
submit_url = "https://api.wavespeed.ai/api/v3/wavespeed-ai/qwen-image/edit-plus-lora"
|
||||
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
|
||||
|
||||
response = await asyncio.to_thread(requests.post, submit_url, headers=headers, data=json.dumps(payload))
|
||||
if response.status_code != 200:
|
||||
raise ConnectionError(f"Failed to submit task. HTTP {response.status_code}: {response.text}")
|
||||
|
||||
request_id = response.json()["data"]["id"]
|
||||
print(f"[WaveSpeedAI] Task submitted successfully. Request ID: {request_id}")
|
||||
|
||||
result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
|
||||
|
||||
output_url = ""
|
||||
final_result = {}
|
||||
while True:
|
||||
print("[WaveSpeedAI] Polling for result...")
|
||||
response = await asyncio.to_thread(requests.get, result_url, headers={"Authorization": f"Bearer {api_key}"})
|
||||
if response.status_code == 200:
|
||||
result = response.json()["data"]
|
||||
status = result["status"]
|
||||
if status == "completed":
|
||||
print("[WaveSpeedAI] Task completed.")
|
||||
output_url = result["outputs"][0]
|
||||
final_result = result
|
||||
break
|
||||
elif status == "failed":
|
||||
raise RuntimeError(f"Task failed: {result.get('error', 'Unknown error')}")
|
||||
else:
|
||||
print(f"[WaveSpeedAI] Status: {status}. Waiting...")
|
||||
else:
|
||||
raise ConnectionError(f"Failed to poll for results. HTTP {response.status_code}: {response.text}")
|
||||
await asyncio.sleep(2)
|
||||
|
||||
print(f"[WaveSpeedAI] Downloading image from: {output_url}")
|
||||
resp = await asyncio.to_thread(requests.get, output_url)
|
||||
if resp.status_code != 200:
|
||||
raise ConnectionError(f"Failed to download image. HTTP {resp.status_code}")
|
||||
|
||||
generation_info = {
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"model": "qwen-image/edit-plus-lora",
|
||||
"parameters": {k: v for k, v in payload.items() if k != 'images'},
|
||||
"input_images_count": len(image_urls),
|
||||
"wavespeed_result": final_result
|
||||
}
|
||||
|
||||
temp_filename = self.get_temp_filename(output_format)
|
||||
image_path, metadata_path = await asyncio.to_thread(
|
||||
self.save_image_and_metadata, resp.content, generation_info, temp_filename
|
||||
)
|
||||
print(f"[WaveSpeedAI] Saved temp image -> {image_path}")
|
||||
print(f"[WaveSpeedAI] Saved metadata -> {metadata_path}")
|
||||
|
||||
pil_image = Image.open(BytesIO(resp.content))
|
||||
output_tensor = self.pil_to_tensor(pil_image)
|
||||
|
||||
return (output_tensor, image_path, json.dumps(generation_info, indent=2))
|
||||
|
||||
except Exception as e:
|
||||
error_info = {
|
||||
"error": f"WaveSpeedAI generation failed: {str(e)}",
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
}
|
||||
print(f"[WaveSpeedAI] Error: {str(e)}")
|
||||
return (None, "", json.dumps(error_info, indent=2))
|
||||
|
||||
def tensor_to_data_uri(self, tensor, image_format="PNG"):
|
||||
pil_img = self.tensor_to_pil(tensor)
|
||||
buffered = BytesIO()
|
||||
pil_img.save(buffered, format=image_format)
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
||||
return f"data:image/{image_format.lower()};base64,{img_str}"
|
||||
|
||||
def tensor_to_pil(self, tensor):
|
||||
if len(tensor.shape) == 4:
|
||||
tensor = tensor[0]
|
||||
arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype("uint8")
|
||||
return Image.fromarray(arr)
|
||||
|
||||
def pil_to_tensor(self, pil_img):
|
||||
img_array = np.array(pil_img).astype(np.float32) / 255.0
|
||||
return torch.from_numpy(img_array)[None,]
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("NaN")
|
||||
|
||||
# Register with ComfyUI
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"WaveSpeedAI_Image": WaveSpeedAI_Image
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"WaveSpeedAI_Image": "WaveSpeedAI Image Generation"
|
||||
}
|
||||
|
||||
__all__ = ["WaveSpeedAI_Image"]
|
||||
Reference in New Issue
Block a user