comfyui-utils-nodes/py/node_volcano.py

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": ""}),
                "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"
}