433 lines
15 KiB
Python
433 lines
15 KiB
Python
import torch
|
|
import numpy as np
|
|
import base64
|
|
import io
|
|
from PIL import Image
|
|
import os
|
|
import requests
|
|
from typing import List, Optional
|
|
|
|
import folder_paths
|
|
import logging
|
|
import yaml
|
|
from .utils import tensor2pil, pil2tensor
|
|
|
|
config_dir = os.path.join(folder_paths.base_path, "config")
|
|
if not os.path.exists(config_dir):
|
|
os.makedirs(config_dir)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class BaseNode:
|
|
def save_config(self, config, config_filename):
|
|
config_path = os.path.join(config_dir, config_filename)
|
|
with open(config_path, 'w', encoding='utf-8') as f:
|
|
yaml.dump(config, f, indent=4)
|
|
|
|
def load_config(self, config_filename):
|
|
config_path = os.path.join(config_dir, config_filename)
|
|
try:
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
config = yaml.load(f, Loader=yaml.FullLoader)
|
|
return config
|
|
except FileNotFoundError:
|
|
return {}
|
|
|
|
def pil_to_base64(self, pil_image):
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format='PNG')
|
|
img_data = buffer.getvalue()
|
|
return base64.b64encode(img_data).decode('utf-8')
|
|
|
|
def encode_image_to_base64(self, image_path):
|
|
with open(image_path, 'rb') as f:
|
|
img_data = f.read()
|
|
return base64.b64encode(img_data).decode('utf-8')
|
|
|
|
def download_image(self, image_url, target_image_name, project_path, timeout=120):
|
|
response = requests.get(image_url, timeout=timeout)
|
|
response.raise_for_status()
|
|
|
|
output_dir = os.path.join(folder_paths.get_output_directory(), project_path)
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
|
|
image_path = os.path.join(output_dir, target_image_name)
|
|
with open(image_path, 'wb') as f:
|
|
f.write(response.content)
|
|
|
|
return image_path
|
|
|
|
class VolcanoOfficialNode(BaseNode):
|
|
def __init__(self):
|
|
from volcengine.visual.VisualService import VisualService
|
|
self.visual_service = VisualService()
|
|
|
|
def save_config(self, config):
|
|
super().save_config(config, 'volcano_config.yml')
|
|
|
|
def load_config(self):
|
|
return super().load_config('volcano_config.yml')
|
|
|
|
def setup_credentials(self, ak, sk):
|
|
"""设置API凭证"""
|
|
if ak and sk:
|
|
self.save_config({"ak": ak, "sk": sk})
|
|
else:
|
|
config = self.load_config()
|
|
ak = config.get("ak")
|
|
sk = config.get("sk")
|
|
if not ak or not sk:
|
|
raise Exception("volcano engine ak or sk not found")
|
|
|
|
self.visual_service.set_ak(ak)
|
|
self.visual_service.set_sk(sk)
|
|
return ak, sk
|
|
|
|
def request_volcano_api(self, access_key, secret_key, form_data):
|
|
"""通用的火山引擎API请求方法"""
|
|
# 设置凭证
|
|
self.setup_credentials(access_key, secret_key)
|
|
|
|
# 发送请求，如果失败则重试一次
|
|
max_retries = 2
|
|
last_exception = None
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
resp = self.visual_service.cv_process(form_data)
|
|
# 检查响应状态
|
|
if resp.get('code') == 50429:
|
|
if attempt < max_retries - 1:
|
|
logger.warning(f"volcano api error, retry... ({attempt + 1}/{max_retries})")
|
|
continue
|
|
else:
|
|
raise Exception(f"volcano api error, retry failed: {resp}")
|
|
|
|
# 解码返回的图像
|
|
img_base64 = resp['data']['binary_data_base64'][0]
|
|
img_data = base64.b64decode(img_base64)
|
|
result_image = Image.open(io.BytesIO(img_data))
|
|
|
|
# 删除图片base64，方便print
|
|
resp['data']['binary_data_base64'][0] = ""
|
|
logger.debug(f"volcano api response: {resp}")
|
|
|
|
return result_image, resp
|
|
|
|
except Exception as e:
|
|
last_exception = e
|
|
error_msg = str(e)
|
|
|
|
# 检查是否超qps
|
|
if ('"code":50429' in error_msg):
|
|
if attempt < max_retries - 1:
|
|
logger.warning(f"volcano api error, retry... ({attempt + 1}/{max_retries}): {error_msg}")
|
|
continue
|
|
|
|
raise e
|
|
|
|
raise last_exception
|
|
|
|
class ArkBaseNode(BaseNode):
|
|
def save_config(self, config):
|
|
super().save_config(config, 'ark_config.yml')
|
|
|
|
def load_config(self):
|
|
return super().load_config('ark_config.yml')
|
|
|
|
def calculate_image_size(self, level, aspect_ratio):
|
|
size_2k = {
|
|
"1:1": (2048, 2048),
|
|
"4:3": (2304, 1728),
|
|
"3:4": (1728, 2304),
|
|
"16:9": (2560, 1440),
|
|
"9:16": (1440, 2560),
|
|
"3:2": (2496, 1664),
|
|
"2:3": (1664, 2496),
|
|
"21:9": (3024, 1296)
|
|
}
|
|
|
|
if aspect_ratio not in size_2k:
|
|
raise Exception(f"不支持的宽高比: {aspect_ratio}")
|
|
|
|
width_2k, height_2k = size_2k[aspect_ratio]
|
|
|
|
if level == "1k":
|
|
width = width_2k // 2
|
|
height = height_2k // 2
|
|
elif level == "2k":
|
|
width = width_2k
|
|
height = height_2k
|
|
elif level == "4k":
|
|
width = width_2k * 2
|
|
height = height_2k * 2
|
|
else:
|
|
raise Exception(f"不支持的等级: {level}")
|
|
|
|
return f"{width}x{height}"
|
|
|
|
def setup_credentials(self, api_key, api_url=None, model=None):
|
|
config = self.load_config()
|
|
if api_key:
|
|
config = {"api_key": api_key}
|
|
if api_url:
|
|
config["api_url"] = api_url
|
|
if model:
|
|
config["model"] = model
|
|
self.save_config(config)
|
|
else:
|
|
api_key = config.get("api_key")
|
|
if not api_key:
|
|
raise Exception("ark api key not found")
|
|
|
|
api_url = config.get("api_url", "https://ark.ap-southeast.bytepluses.com/api/v3/images/generations") if api_url is None else api_url
|
|
model = config.get("model", "seedream-4-0-250828") if model is None else model
|
|
|
|
return api_key, api_url, model
|
|
|
|
def request_ark_api(self, prompt, image_paths=None, max_images=1, image_size="1440x2560",
|
|
api_key=None, api_url=None, model=None, timeout=None):
|
|
api_key, api_url, model = self.setup_credentials(api_key, api_url, model)
|
|
|
|
images = []
|
|
if image_paths and len(image_paths) > 0:
|
|
for image_path in image_paths:
|
|
if image_path.startswith("http://") or image_path.startswith("https://"):
|
|
images.append(image_path)
|
|
else:
|
|
images.append(f"data:image/png;base64,{self.encode_image_to_base64(image_path)}")
|
|
|
|
request_data = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"image": images,
|
|
"sequential_image_generation": "auto",
|
|
"sequential_image_generation_options": {
|
|
"max_images": max_images
|
|
},
|
|
"response_format": "url",
|
|
"size": image_size,
|
|
"stream": False,
|
|
"watermark": False
|
|
}
|
|
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}"
|
|
}
|
|
|
|
logger.info(f"ark请求数据 image_size: {image_size}, prompt: {prompt}")
|
|
response = requests.post(api_url, headers=headers, json=request_data, timeout=timeout)
|
|
response.raise_for_status()
|
|
|
|
result_data = response.json()
|
|
logger.info(f"ark请求成功，生成{len(result_data.get('data', []))}张图片")
|
|
|
|
image_paths = []
|
|
for i, image_data in enumerate(result_data.get('data', [])):
|
|
image_url = image_data.get('url')
|
|
if image_url:
|
|
image_path = self.download_image(image_url, f"ark_result_{i}.png", "ark", timeout)
|
|
image_paths.append(image_path)
|
|
|
|
return image_paths
|
|
|
|
class VolcanoOutpaintingNode(VolcanoOfficialNode):
|
|
"""火山引擎图像扩展节点"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
"access_key": ("STRING", {"default": ""}),
|
|
"secret_key": ("STRING", {"default": ""}),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 999999})
|
|
},
|
|
"optional": {
|
|
"version": ("STRING", {"default": "i2i_outpainting", "tooltip": "this value is the reqKey of the volcano api."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("result_image",)
|
|
FUNCTION = "process_outpainting"
|
|
CATEGORY = "image/volcano"
|
|
|
|
def process_outpainting(self, image, mask, access_key, secret_key, seed, version="i2i_outpainting"):
|
|
try:
|
|
# 使用节点库的转换函数
|
|
pil_image = tensor2pil(image)
|
|
pil_mask = tensor2pil(mask)
|
|
|
|
# 确保mask是灰度图像
|
|
if pil_mask.mode != 'L':
|
|
pil_mask = pil_mask.convert('L')
|
|
|
|
# 转换为base64
|
|
image_base64 = self.pil_to_base64(pil_image)
|
|
mask_base64 = self.pil_to_base64(pil_mask)
|
|
|
|
# 构建请求参数
|
|
form_data = {
|
|
"req_key": version,
|
|
"binary_data_base64": [image_base64, mask_base64],
|
|
"top": 0,
|
|
"left": 0,
|
|
"bottom": 0,
|
|
"right": 0,
|
|
"seed": seed
|
|
}
|
|
|
|
# 调用基类的通用API请求方法
|
|
result_image, _ = self.request_volcano_api(access_key, secret_key, form_data)
|
|
|
|
# 使用节点库的转换函数转换结果为tensor
|
|
result_tensor = pil2tensor(result_image)
|
|
|
|
return (result_tensor,)
|
|
|
|
except Exception as e:
|
|
logger.exception(e)
|
|
raise e
|
|
|
|
class VolcanoImageEditNode(VolcanoOfficialNode):
|
|
"""火山引擎图像编辑节点"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"access_key": ("STRING", {"default": ""}),
|
|
"secret_key": ("STRING", {"default": ""}),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}),
|
|
"scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.01})
|
|
},
|
|
"optional": {
|
|
"version": (["byteedit_v2.0","seededit_v3.0"], {"default": "seededit_v3.0"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("edited_image",)
|
|
FUNCTION = "process_image_edit"
|
|
CATEGORY = "image/volcano"
|
|
|
|
def process_image_edit(self, image, prompt, access_key, secret_key, seed, scale=0.5, version="seededit_v3.0"):
|
|
try:
|
|
# 使用节点库的转换函数
|
|
pil_image = tensor2pil(image)
|
|
|
|
# 转换为base64
|
|
image_base64 = self.pil_to_base64(pil_image)
|
|
binary_data = [image_base64]
|
|
# 构建请求参数
|
|
form_data = {
|
|
"req_key": version,
|
|
"binary_data_base64": binary_data,
|
|
"prompt": prompt,
|
|
"seed": seed,
|
|
"scale": scale
|
|
}
|
|
|
|
# 调用基类的通用API请求方法
|
|
result_image, _ = self.request_volcano_api(access_key, secret_key, form_data)
|
|
|
|
# 使用节点库的转换函数转换结果为tensor
|
|
result_tensor = pil2tensor(result_image)
|
|
|
|
return (result_tensor,)
|
|
|
|
except Exception as e:
|
|
logger.exception(e)
|
|
raise e
|
|
|
|
class VolcanoArkImageEditNode(ArkBaseNode):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"level": (["1k", "2k", "4k"], {"default": "2k"}),
|
|
"aspect_ratio": (["1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "21:9"], {"default": "9:16"}),
|
|
"max_images": ("INT", {"default": 1, "min": 1, "max": 10}),
|
|
"timeout": ("INT", {"default": 120, "min": 10, "max": 600})
|
|
},
|
|
"optional": {
|
|
"image2": ("IMAGE",),
|
|
"image3": ("IMAGE",),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
"api_url": ("STRING", {"default": "", "multiline": True}),
|
|
"model": ("STRING", {"default": ""})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("edited_image",)
|
|
FUNCTION = "process_image_edit"
|
|
CATEGORY = "image/ark"
|
|
|
|
def process_image_edit(self, image, prompt, level, aspect_ratio, max_images, timeout,
|
|
image2=None, image3=None, api_key="", api_url="", model=""):
|
|
pil_image = tensor2pil(image)
|
|
|
|
temp_dir = os.path.join(folder_paths.get_temp_directory(), "ark")
|
|
if not os.path.exists(temp_dir):
|
|
os.makedirs(temp_dir)
|
|
|
|
temp_image_path = os.path.join(temp_dir, "input_image.png")
|
|
pil_image.save(temp_image_path)
|
|
|
|
image_paths = [temp_image_path]
|
|
|
|
if image2 is not None:
|
|
pil_image2 = tensor2pil(image2)
|
|
temp_image_path2 = os.path.join(temp_dir, "input_image2.png")
|
|
pil_image2.save(temp_image_path2)
|
|
image_paths.append(temp_image_path2)
|
|
|
|
if image3 is not None:
|
|
pil_image3 = tensor2pil(image3)
|
|
temp_image_path3 = os.path.join(temp_dir, "input_image3.png")
|
|
pil_image3.save(temp_image_path3)
|
|
image_paths.append(temp_image_path3)
|
|
|
|
image_size = self.calculate_image_size(level, aspect_ratio)
|
|
|
|
image_paths_result = self.request_ark_api(
|
|
prompt=prompt,
|
|
image_paths=image_paths,
|
|
max_images=max_images,
|
|
image_size=image_size,
|
|
api_key=api_key if api_key else None,
|
|
api_url=api_url if api_url else None,
|
|
model=model if model else None,
|
|
timeout=timeout
|
|
)
|
|
|
|
if not image_paths_result or len(image_paths_result) == 0:
|
|
raise Exception("ark api returned no images")
|
|
|
|
result_image = Image.open(image_paths_result[0])
|
|
result_tensor = pil2tensor(result_image)
|
|
|
|
return (result_tensor,)
|
|
|
|
# 节点映射
|
|
NODE_CLASS_MAPPINGS = {
|
|
"VolcanoOutpaintingNode": VolcanoOutpaintingNode,
|
|
"VolcanoImageEditNode": VolcanoImageEditNode,
|
|
"VolcanoArkImageEditNode": VolcanoArkImageEditNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"VolcanoOutpaintingNode": "volcano outpainting",
|
|
"VolcanoImageEditNode": "volcano image edit",
|
|
"VolcanoArkImageEditNode": "volcano ark image edit"
|
|
}
|