90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
import io
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from . import veida_client as kc
|
|
|
|
RATIOS = ["1:1", "16:9", "9:16", "4:3", "3:4"]
|
|
|
|
|
|
def _to_tensor(data: bytes) -> torch.Tensor:
|
|
img = Image.open(io.BytesIO(data)).convert("RGB")
|
|
return torch.from_numpy(np.asarray(img).astype(np.float32) / 255.0)[None, ...]
|
|
|
|
|
|
def _to_png(image: torch.Tensor) -> bytes:
|
|
arr = (image[0].cpu().numpy().clip(0, 1) * 255).astype(np.uint8)
|
|
buf = io.BytesIO()
|
|
Image.fromarray(arr).save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
KEY_TOOLTIP = (
|
|
"Optional. Leave empty to use the free tier (no account, watermarked, a few images a day). "
|
|
"Paste a key from veida.ai/settings/apikeys to run on your account with any image model "
|
|
"your plan includes. VEIDA_API_KEY in the environment works too."
|
|
)
|
|
|
|
|
|
class VeidaGenerateImage:
|
|
CATEGORY = "Veida"
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "url")
|
|
FUNCTION = "generate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": "matte black ceramic mug on pale oak, soft window light"}),
|
|
"aspect_ratio": (RATIOS, {"default": "1:1"}),
|
|
},
|
|
"optional": {
|
|
"api_key": ("STRING", {"default": "", "tooltip": KEY_TOOLTIP}),
|
|
"model": ("STRING", {"default": "", "tooltip": "Only used with an API key, e.g. gpt-image-2, gpt-image-2-5, nano-banana-2, seedream-5-lite"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF, "tooltip": "Change to force a new run"}),
|
|
},
|
|
}
|
|
|
|
def generate(self, prompt, aspect_ratio, api_key="", model="", seed=0):
|
|
url = kc.run(prompt, aspect_ratio=aspect_ratio, api_key=api_key, model=model)
|
|
return (_to_tensor(kc.download(url)), url)
|
|
|
|
|
|
class VeidaEditImage:
|
|
CATEGORY = "Veida"
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "url")
|
|
FUNCTION = "edit"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"instruction": ("STRING", {"multiline": True, "default": "put the product on a white marble surface; keep the label unchanged"}),
|
|
},
|
|
"optional": {
|
|
"api_key": ("STRING", {"default": "", "tooltip": KEY_TOOLTIP}),
|
|
"model": ("STRING", {"default": "", "tooltip": "Only used with an API key, e.g. gpt-image-2-5, nano-banana-2"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF, "tooltip": "Change to force a new run"}),
|
|
},
|
|
}
|
|
|
|
def edit(self, image, instruction, api_key="", model="", seed=0):
|
|
source = kc.upload_png(_to_png(image))
|
|
url = kc.run(instruction, image_url=source, api_key=api_key, model=model)
|
|
return (_to_tensor(kc.download(url)), url)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"VeidaGenerateImage": VeidaGenerateImage,
|
|
"VeidaEditImage": VeidaEditImage,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"VeidaGenerateImage": "Veida Generate Image",
|
|
"VeidaEditImage": "Veida Edit Image",
|
|
}
|