237 lines
9.7 KiB
Python
237 lines
9.7 KiB
Python
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", "STRING",)
|
|
RETURN_NAMES = ("image", "image_path", "generation_info", "generation_time",)
|
|
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...")
|
|
start_time = time.time()
|
|
|
|
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-2511-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)
|
|
|
|
generation_time = self.format_generation_time(time.time() - start_time)
|
|
print(f"[WaveSpeedAI] Generation completed in {generation_time} seconds")
|
|
return (output_tensor, image_path, json.dumps(generation_info, indent=2), generation_time)
|
|
|
|
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)}")
|
|
generation_time = self.format_generation_time(time.time() - start_time)
|
|
return (None, "", json.dumps(error_info, indent=2), generation_time)
|
|
|
|
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,]
|
|
|
|
def format_generation_time(self, elapsed_seconds):
|
|
"""Format elapsed time as 'seconds-centiseconds' (e.g., '14-50' for 14.50 seconds)"""
|
|
seconds = int(elapsed_seconds)
|
|
centiseconds = int((elapsed_seconds - seconds) * 100)
|
|
return f"{seconds}-{centiseconds:02d}"
|
|
|
|
@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"]
|