177 lines
6.7 KiB
Python
177 lines
6.7 KiB
Python
import base64
|
|
import io
|
|
|
|
import numpy as np
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from .aimlapi_shared import aspect_ratio_to_size, validate_api_key
|
|
|
|
|
|
class AIMLAPIImageGenerator:
|
|
"""
|
|
Text-to-Image generation using Google Nano Banana Pro or ByteDance Seedream 4.5.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"api_key": ("STRING", {
|
|
"multiline": False,
|
|
"default": "",
|
|
"password": True,
|
|
"tooltip": "Your AI/ML API key from https://aimlapi.com/app/keys",
|
|
}),
|
|
"model": (["google/nano-banana-pro", "bytedance/seedream-4-5"], {
|
|
"default": "google/nano-banana-pro",
|
|
"tooltip": "Select the image generation model",
|
|
}),
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "A serene landscape with mountains and a lake at sunset",
|
|
"tooltip": "Text description of the image to generate",
|
|
}),
|
|
},
|
|
"optional": {
|
|
"aspect_ratio": (["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "21:9", "5:4", "4:5"], {
|
|
"default": "1:1",
|
|
"tooltip": "Aspect ratio of the generated image",
|
|
}),
|
|
"resolution": (["1K", "2K", "4K"], {
|
|
"default": "1K",
|
|
"tooltip": "Image resolution",
|
|
}),
|
|
"num_images": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 4,
|
|
"step": 1,
|
|
"tooltip": "Number of images to generate (Nano Banana only, 1-4)",
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": -1,
|
|
"min": -1,
|
|
"max": 2147483647,
|
|
"tooltip": "Seed for reproducibility (Seedream only, -1 for random)",
|
|
}),
|
|
"guidance_scale": ("FLOAT", {
|
|
"default": 3.5,
|
|
"min": 0.0,
|
|
"max": 20.0,
|
|
"step": 0.1,
|
|
"tooltip": "Guidance scale (Seedream only)",
|
|
}),
|
|
"watermark": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Add invisible watermark (Seedream only)",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "generate_image"
|
|
CATEGORY = "AI/ML API"
|
|
|
|
def generate_image(self, api_key, model, prompt, aspect_ratio="1:1",
|
|
resolution="1K", num_images=1, seed=-1, guidance_scale=3.5, watermark=False):
|
|
api_key = validate_api_key(api_key)
|
|
|
|
url = "https://api.aimlapi.com/v1/images/generations"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
payload = {
|
|
"prompt": prompt,
|
|
"model": model,
|
|
"response_format": "url",
|
|
}
|
|
|
|
if model == "bytedance/seedream-4-5":
|
|
payload["size"] = aspect_ratio_to_size(aspect_ratio, resolution)
|
|
payload["watermark"] = watermark
|
|
payload["guidance_scale"] = guidance_scale
|
|
if seed != -1:
|
|
payload["seed"] = seed
|
|
elif model == "google/nano-banana-pro":
|
|
payload["aspect_ratio"] = aspect_ratio
|
|
payload["resolution"] = resolution
|
|
payload["num_images"] = num_images
|
|
|
|
try:
|
|
response = requests.post(url, headers=headers, json=payload, timeout=120)
|
|
if response.status_code != 200:
|
|
self._raise_image_error(response)
|
|
|
|
result = response.json()
|
|
image_data_list = None
|
|
if "data" in result and result["data"]:
|
|
image_data_list = result["data"]
|
|
elif "images" in result and result["images"]:
|
|
image_data_list = result["images"]
|
|
|
|
if not image_data_list:
|
|
raise ValueError(f"No images returned from API. Response: {result}")
|
|
|
|
images = []
|
|
for img_data in image_data_list:
|
|
image = self._download_image(img_data)
|
|
image_np = np.array(image).astype(np.float32) / 255.0
|
|
image_tensor = torch.from_numpy(image_np)[None,]
|
|
images.append(image_tensor)
|
|
|
|
output = torch.cat(images, dim=0) if len(images) > 1 else images[0]
|
|
return (output,)
|
|
|
|
except requests.exceptions.RequestException as exc:
|
|
raise ValueError(f"Network error: {exc}")
|
|
|
|
@staticmethod
|
|
def _raise_image_error(response):
|
|
error_msg = f"API request failed with status {response.status_code}"
|
|
try:
|
|
error_data = response.json()
|
|
if "error" in error_data:
|
|
error_msg += f": {error_data['error']}"
|
|
elif "message" in error_data:
|
|
error_msg += f": {error_data['message']}"
|
|
else:
|
|
error_msg += f": {error_data}"
|
|
except Exception:
|
|
error_msg += f": {response.text[:300]}"
|
|
|
|
if response.status_code == 403:
|
|
error_msg += "\n\nTroubleshooting 403 Forbidden:\n"
|
|
error_msg += "- Verify your API key is correct at https://aimlapi.com/app/keys\n"
|
|
error_msg += "- Check if your API key has sufficient credits\n"
|
|
error_msg += "- Ensure your API key hasn't expired\n"
|
|
error_msg += "- Try generating a new API key"
|
|
elif response.status_code == 401:
|
|
error_msg += "\n\nAPI key is invalid or missing. Get a valid key from https://aimlapi.com/app/keys"
|
|
elif response.status_code == 429:
|
|
error_msg += "\n\nRate limit exceeded. Please wait a moment and try again."
|
|
elif response.status_code >= 500:
|
|
error_msg += "\n\nAPI server error. Please try again later."
|
|
|
|
raise ValueError(error_msg)
|
|
|
|
@staticmethod
|
|
def _download_image(img_data):
|
|
if "url" in img_data:
|
|
img_response = requests.get(img_data["url"], timeout=30)
|
|
img_response.raise_for_status()
|
|
image = Image.open(io.BytesIO(img_response.content))
|
|
elif "b64_json" in img_data and img_data["b64_json"]:
|
|
image_bytes = base64.b64decode(img_data["b64_json"])
|
|
image = Image.open(io.BytesIO(image_bytes))
|
|
else:
|
|
raise ValueError(f"No image URL or base64 data in response: {img_data}")
|
|
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
return image
|