249 lines
8.7 KiB
Python
249 lines
8.7 KiB
Python
import json
|
|
import sys
|
|
import os
|
|
import io
|
|
import base64
|
|
import datetime
|
|
import hashlib
|
|
import hmac
|
|
import requests
|
|
import time
|
|
from PIL import Image
|
|
import torch
|
|
from .imagefunc import tensor2pil, pil2tensor, log, get_api_key, fit_resize_image
|
|
|
|
jimeng_i2i_model_list = ["jimeng_i2i_v30"]
|
|
jimeng_t2i_model_list = ["jimeng_high_aes_general_v21_L"]
|
|
|
|
|
|
method = 'POST'
|
|
host = 'visual.volcengineapi.com'
|
|
region = 'cn-north-1'
|
|
endpoint = 'https://visual.volcengineapi.com'
|
|
service = 'cv'
|
|
image_max_side_length = 4096
|
|
|
|
def sign(key, msg):
|
|
return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()
|
|
|
|
def getSignatureKey(key, dateStamp, regionName, serviceName):
|
|
kDate = sign(key.encode('utf-8'), dateStamp)
|
|
kRegion = sign(kDate, regionName)
|
|
kService = sign(kRegion, serviceName)
|
|
kSigning = sign(kService, 'request')
|
|
return kSigning
|
|
|
|
def formatQuery(parameters):
|
|
request_parameters_init = ''
|
|
for key in sorted(parameters):
|
|
request_parameters_init += key + '=' + parameters[key] + '&'
|
|
request_parameters = request_parameters_init[:-1]
|
|
return request_parameters
|
|
|
|
|
|
def signV4Request(access_key, secret_key, service, req_query, req_body):
|
|
if access_key is None or secret_key is None:
|
|
print('No access key is available.')
|
|
sys.exit()
|
|
|
|
t = datetime.datetime.utcnow()
|
|
current_date = t.strftime('%Y%m%dT%H%M%SZ')
|
|
# current_date = '20210818T095729Z'
|
|
datestamp = t.strftime('%Y%m%d') # Date w/o time, used in credential scope
|
|
canonical_uri = '/'
|
|
canonical_querystring = req_query
|
|
signed_headers = 'content-type;host;x-content-sha256;x-date'
|
|
payload_hash = hashlib.sha256(req_body.encode('utf-8')).hexdigest()
|
|
content_type = 'application/json'
|
|
canonical_headers = 'content-type:' + content_type + '\n' + 'host:' + host + \
|
|
'\n' + 'x-content-sha256:' + payload_hash + \
|
|
'\n' + 'x-date:' + current_date + '\n'
|
|
canonical_request = method + '\n' + canonical_uri + '\n' + canonical_querystring + \
|
|
'\n' + canonical_headers + '\n' + signed_headers + '\n' + payload_hash
|
|
# print(canonical_request)
|
|
algorithm = 'HMAC-SHA256'
|
|
credential_scope = datestamp + '/' + region + '/' + service + '/' + 'request'
|
|
string_to_sign = algorithm + '\n' + current_date + '\n' + credential_scope + '\n' + hashlib.sha256(
|
|
canonical_request.encode('utf-8')).hexdigest()
|
|
# print(string_to_sign)
|
|
signing_key = getSignatureKey(secret_key, datestamp, region, service)
|
|
# print(signing_key)
|
|
signature = hmac.new(signing_key, (string_to_sign).encode(
|
|
'utf-8'), hashlib.sha256).hexdigest()
|
|
# print(signature)
|
|
|
|
authorization_header = algorithm + ' ' + 'Credential=' + access_key + '/' + \
|
|
credential_scope + ', ' + 'SignedHeaders=' + \
|
|
signed_headers + ', ' + 'Signature=' + signature
|
|
# print(authorization_header)
|
|
headers = {'X-Date': current_date,
|
|
'Authorization': authorization_header,
|
|
'X-Content-Sha256': payload_hash,
|
|
'Content-Type': content_type
|
|
}
|
|
# print(headers)
|
|
|
|
# ************* SEND THE REQUEST *************
|
|
request_url = endpoint + '?' + canonical_querystring
|
|
|
|
# print('\nBEGIN REQUEST++++++++++++++++++++++++++++++++++++')
|
|
# print('Request URL = ' + request_url)
|
|
try:
|
|
r = requests.post(request_url, headers=headers, data=req_body)
|
|
except Exception as err:
|
|
print(f'error occurred: {err}')
|
|
raise
|
|
else:
|
|
# print('\nRESPONSE++++++++++++++++++++++++++++++++++++')
|
|
# print(f'Response code: {r.status_code}\n')
|
|
# 使用 replace 方法将 \u0026 替换为 &
|
|
resp_str = r.text.replace("\\u0026", "&")
|
|
# print(f'Response body: {resp_str}\n')
|
|
return json.loads(resp_str)
|
|
|
|
def round_to_multiple_of_16(x):
|
|
return int(round(x / 16)) * 16
|
|
|
|
def get_closest_api_size(orig_width, orig_height):
|
|
"""
|
|
根据原始图片尺寸，返回最接近的API尺寸。
|
|
"""
|
|
aspect = orig_width / orig_height
|
|
max_width, max_height = 2016, 1536
|
|
max_aspect = max_width / max_height
|
|
|
|
if aspect > max_aspect:
|
|
# 限制宽度
|
|
width = max_width
|
|
height = width / aspect
|
|
else:
|
|
# 限制高度
|
|
height = max_height
|
|
width = aspect * height
|
|
|
|
# 四舍五入到16的倍数
|
|
width = round_to_multiple_of_16(width)
|
|
height = round_to_multiple_of_16(height)
|
|
|
|
# 再次确保合法范围
|
|
width = min(max(width, 512), 2016)
|
|
height = min(max(height, 512), 1536)
|
|
|
|
return width, height
|
|
|
|
|
|
|
|
class LS_Jimeng_i2i_API:
|
|
|
|
def __init__(self):
|
|
self.NODE_NAME = 'Jimeng Image2Image API'
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model": (jimeng_i2i_model_list,),
|
|
"time_out": ("INT", {"default": 300, "min": 1, "max": 3600, "step": 1}), # 300s = 5min
|
|
"scale": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.1}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 1e18, "step": 1}),
|
|
"prompt": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
"optional": {
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = 'run_jimeng_i2i_api'
|
|
CATEGORY = '😺dzNodes/LayerUtility'
|
|
|
|
def run_jimeng_i2i_api(self, image, model, time_out, scale, seed, prompt):
|
|
ret_images = []
|
|
access_key = get_api_key("volcengine_AccessKeyId")
|
|
secret_key = get_api_key("volcengine_SecretAccessKey")
|
|
|
|
# image shape = b, h, w, c
|
|
orig_width = image.shape[2]
|
|
orig_height = image.shape[1]
|
|
output_width, output_height = get_closest_api_size(orig_width, orig_height)
|
|
|
|
for img in image:
|
|
img = torch.unsqueeze(img, 0)
|
|
orig_image = tensor2pil(img).convert('RGB')
|
|
orig_image = fit_resize_image(orig_image, output_width, output_height, 'fill', Image.BICUBIC)
|
|
buffered = io.BytesIO()
|
|
orig_image.save(buffered, format="JPEG")
|
|
image_bytes = buffered.getvalue()
|
|
image_base64 = base64.b64encode(image_bytes).decode("utf-8")
|
|
body_params = {
|
|
"req_key": model,
|
|
"binary_data_base64": [image_base64],
|
|
"prompt": prompt,
|
|
"seed": seed,
|
|
"scale": scale,
|
|
"width": output_width,
|
|
"height": output_height,
|
|
}
|
|
formatted_body = json.dumps(body_params)
|
|
|
|
query_params = {
|
|
'Action': 'CVSync2AsyncSubmitTask',
|
|
'Version': '2022-08-31',
|
|
}
|
|
formatted_query = formatQuery(query_params)
|
|
|
|
# 发送请求
|
|
|
|
result = signV4Request(access_key, secret_key, service, formatQuery(query_params), formatted_body)
|
|
response_code = result["code"]
|
|
if response_code > 10000:
|
|
print(f"Error: {result}")
|
|
raise Exception(f"error code: {response_code}")
|
|
task_id = result["data"]["task_id"]
|
|
log(f"Send Request to Jimeng API: task_id = {task_id}")
|
|
|
|
# 获取结果
|
|
start_time = time.time()
|
|
check_interval = 1
|
|
check_timeout = 60
|
|
|
|
query_params["Action"] = 'CVSync2AsyncGetResult'
|
|
formatted_query = formatQuery(query_params)
|
|
|
|
body_params = {
|
|
"req_key": "jimeng_i2i_v30",
|
|
"task_id": task_id,
|
|
}
|
|
formatted_body = json.dumps(body_params)
|
|
|
|
while True:
|
|
time.sleep(check_interval)
|
|
result = signV4Request(access_key, secret_key, service, formatted_query, formatted_body)
|
|
query_status = result["data"]["status"]
|
|
if query_status == "done":
|
|
# print(f"query_status: {query_status}")
|
|
break
|
|
if time.time() - start_time > check_timeout:
|
|
raise Exception("Jimeng API Timeout")
|
|
|
|
end_time = time.time()
|
|
use_time = round(end_time - start_time, 2)
|
|
log(f"Jimeng API responded in {use_time} seconds.")
|
|
|
|
base64_str = result["data"]["binary_data_base64"][0]
|
|
image_data = base64.b64decode(base64_str)
|
|
ret_image = Image.open(io.BytesIO(image_data))
|
|
ret_images.append(pil2tensor(ret_image))
|
|
|
|
return (torch.cat(ret_images, dim=0),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LayerUtility: JimengI2IAPI": LS_Jimeng_i2i_API,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LayerUtility: JimengI2IAPI": "LayerUtility: Jimeng Imgae to Image API (Advance)",
|
|
} |