utils.py

import random
import io
import zipfile
import requests
import json
import base64
import math
import gradio as gr
import hashlib
import time

from PIL import Image

jwt_token = ''
url = "https://image.novelai.net/ai/generate-image"
headers = {}

def set_token(token):
    global jwt_token, headers
    if jwt_token == token:
        return
    jwt_token = token
    headers = {
            "Authorization": f"Bearer {jwt_token}",
            "Content-Type": "application/json",
            "Origin": "https://novelai.net",
            "Referer": "https://novelai.net/"
        }

def get_remain_anlas():
    try:
        data = requests.get("https://api.novelai.net/user/data", headers=headers).content
        anlas = json.loads(data)['subscription']['trainingStepsLeft']
        return anlas['fixedTrainingStepsLeft'] + anlas['purchasedTrainingSteps']
    except:
        return '获取失败，err:' + str(data)

def calculate_cost(width, height, steps=28, sm=False, dyn=False, strength=1, rmbg=False):
    pixels = width * height
    if pixels <= 1048576 and steps <= 28 and not rmbg:
        return 0
    dyn = sm and dyn
    L = math.ceil(2951823174884865e-21 * pixels + 5.753298233447344e-7 * pixels * steps)
    L *= 1.4 if dyn else (1.2 if sm else 1)
    L = max(math.ceil(L * strength), 2)
    return L * 3 + 5 if rmbg else L

def generate_novelai_image(
    model="nai-diffusion-3",
    input_text="", 
    negative_prompt="", 
    seed=-1, 
    scale=5.0, 
    width=1024, 
    height=1024, 
    steps=28, 
    sampler="k_euler",
    schedule='native',
    smea=False,
    dyn=False,
    dyn_threshold=False,
    cfg_rescale=0,
    variety=False,
    ref_images=None,
    info_extracts=[],
    ref_strs=[],
    vibe_files=[],
    str_norm=True,
    i2i_image=None,
    i2i_str=0.7,
    i2i_noise=0,
    overlay=True,
    inp_img=None,
    inp_str=1,
    selection='i2i',
    auto_pos=True,
    char_prompts=[],
    char_ucs=[],
    char_coords_x=[],
    char_coords_y=[],
    legacy=False,
    chr_image=None,
    fidelity=1,
    style_aware=True
):
    # Assign a random seed if seed is -1
    if seed == -1:
        seed = random.randint(0, 2 ** 32 - 1)

    characterPrompts = []
    for i in range(len(char_prompts)):
        characterPrompts.append({"prompt": char_prompts[i], "uc": char_ucs[i], "center": {'x': round(char_coords_x[i] * 0.2 - 0.1, 1), 'y': round(char_coords_y[i] * 0.2 - 0.1, 1)}})
    
    # Define the payload
    payload = {
        "action": "generate",
        "input": input_text,
        "model": model,
        "parameters": {
            "width": width,
            "height": height,
            "scale": scale,
            "sampler": sampler,
            "steps": steps,
            "n_samples": 1,
            "ucPreset": 0,
            "cfg_rescale": cfg_rescale,
            "controlnet_strength": 1,
            "dynamic_thresholding": dyn_threshold,
            "params_version": 3,
            "legacy": False,
            "legacy_uc": legacy,
            "legacy_v3_extend": False,
            "negative_prompt": negative_prompt,
            "noise_schedule": schedule,
            "qualityToggle": True,
            "reference_strength_multiple": ref_strs,
            "normalize_reference_strength_multiple": str_norm,
            "seed": seed,
            "skip_cfg_above_sigma": (58 if model.startswith('nai-diffusion-4-5') else 19) if variety else None,
            "sm": smea,
            "sm_dyn": dyn,
            "add_original_image": overlay,
            "characterPrompts": characterPrompts,
            "use_coords": not auto_pos,
            "deliberate_euler_ancestral_bug": False,
            "prefer_brownian": True
        }
    }
    if model.startswith("nai-diffusion-4"):
        payload["parameters"]["v4_prompt"] = {"caption": {"base_caption": input_text, "char_captions": []}, "use_coords": not auto_pos, "use_order": True}
        payload["parameters"]["v4_negative_prompt"] = {"caption": {"base_caption": negative_prompt, "char_captions": []}}
        for i in range(len(char_prompts)):
            payload["parameters"]["v4_prompt"]["caption"]["char_captions"].append({"char_caption": char_prompts[i], "centers": [{'x': round(char_coords_x[i] * 0.2 - 0.1, 1), 'y': round(char_coords_y[i] * 0.2 - 0.1, 1)}]})
            payload["parameters"]["v4_negative_prompt"]["caption"]["char_captions"].append({"char_caption": char_ucs[i], "centers": [{'x': round(char_coords_x[i] * 0.2 - 0.1, 1), 'y': round(char_coords_y[i] * 0.2 - 0.1, 1)}]})
    if ref_images != None:
        payload['parameters']['reference_image_multiple'] = [image2base64(image[0]) for image in ref_images]
        payload['parameters']['reference_information_extracted_multiple'] = info_extracts
    if vibe_files != None and model.startswith('nai-diffusion-4'):
        vibes = []
        for v in vibe_files:
            with open(v) as f:
                data = json.load(f)['encodings']
                vibes.append(data.popitem()[1].popitem()[1]['encoding'])
        payload['parameters']['reference_image_multiple'] = vibes
    if selection == 'inp' and inp_img['background'] != None:
        payload['action'] = "infill"
        payload['model'] = model.replace('-preview', '') + '-inpainting'
        mask = inp_img['layers'][0].resize((width, height))
        mask = mask.resize((mask.size[0] // 8, mask.size[1] // 8), resample=Image.NEAREST).resize(mask.size, resample=Image.NEAREST).point(lambda x: 0 if x < 255 else 255)
        payload['parameters']['mask'] = image2base64(mask)
        payload['parameters']['image'] = image2base64(inp_img['background'])
        payload['parameters']['extra_noise_seed'] = seed
        payload['parameters']['inpaintImg2ImgStrength'] = inp_str
        payload['parameters']['img2img'] = {'strength': inp_str, 'color_correct': True}
    if i2i_image != None and selection == 'i2i':
        payload['action'] = "img2img"
        payload['parameters']['image'] = image2base64(i2i_image)
        payload['parameters']['strength'] = i2i_str
        payload['parameters']['extra_noise_seed'] = seed
        payload["parameters"]['noise'] = i2i_noise
    if chr_image != None:
        payload['parameters']['director_reference_images'] = [image2base64(resize_and_pad_image(chr_image))]
        payload['parameters']['director_reference_descriptions'] = [{'caption': {'base_caption': 'character&style' if style_aware else 'character', 'char_captions': []}, 'legacy_uc': False}]
        payload['parameters']['director_reference_secondary_strength_values'] = [1 - fidelity]
        payload['parameters']['director_reference_strength_values'] = [1]
        payload['parameters']['director_reference_information_extracted'] = [1]
    # Send the POST request
    try:
        response = requests.post(url, json=payload, headers=headers, timeout=180)
    except:
        raise gr.Error('NAI response timeout')

    # Process the response
    if response.headers.get('Content-Type') == 'binary/octet-stream':
        zipfile_in_memory = io.BytesIO(response.content)
        with zipfile.ZipFile(zipfile_in_memory, 'r') as zip_ref:
            file_names = zip_ref.namelist()
            if file_names:
                with zip_ref.open(file_names[0]) as file:
                    return file.read(), payload
            else:
                messages = json.loads(response.content)
                raise gr.Error(str(messages["statusCode"]) + ": " + messages["message"])
    else:
        messages = json.loads(response.content)
        raise gr.Error(str(messages["statusCode"]) + ": " + messages["message"])

def image_from_bytes(data):
    img_file = io.BytesIO(data)
    img_file.seek(0)
    return Image.open(img_file)

def image2base64(img):
    output_buffer = io.BytesIO()
    img.save(output_buffer, format='PNG' if img.mode=='RGBA' else 'JPEG')
    byte_data = output_buffer.getvalue()
    base64_str = base64.b64encode(byte_data).decode()
    return base64_str

def base642image(b64):
    image_bytes = base64.b64decode(b64)
    read_buffer = io.BytesIO(image_bytes)
    img = Image.open(read_buffer)
    return img

def resize_and_pad_image(chr_image):
    """
    根据图像宽高比选择最接近的尺寸，等比缩放并用黑边填充
    
    Args:
        chr_image: PIL Image对象
    
    Returns:
        处理后的PIL Image对象
    """
    # 定义三种目标尺寸
    target_sizes = [
        (1024, 1536),  # 竖版
        (1472, 1472),  # 正方形
        (1536, 1024)   # 横版
    ]
    
    # 获取原图尺寸
    original_width, original_height = chr_image.size
    original_ratio = original_width / original_height
    
    # 计算每个目标尺寸的宽高比，并找出最接近的
    min_diff = float('inf')
    best_size = None
    
    for size in target_sizes:
        target_ratio = size[0] / size[1]
        diff = abs(original_ratio - target_ratio)
        if diff < min_diff:
            min_diff = diff
            best_size = size
    
    # 目标尺寸
    target_width, target_height = best_size
    
    # 计算缩放比例（保持宽高比）
    scale_x = target_width / original_width
    scale_y = target_height / original_height
    scale = min(scale_x, scale_y)  # 选择较小的缩放比例以确保图像完全显示
    
    # 计算缩放后的尺寸
    new_width = int(original_width * scale)
    new_height = int(original_height * scale)
    
    # 缩放图像
    resized_image = chr_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # 创建目标尺寸的黑色背景
    padded_image = Image.new('RGBA', (target_width, target_height), color='black')
    
    # 计算居中位置
    x_offset = (target_width - new_width) // 2
    y_offset = (target_height - new_height) // 2
    
    # 将缩放后的图像粘贴到黑色背景上
    padded_image.paste(resized_image, (x_offset, y_offset))
    
    return padded_image

def augment_image(image, width, height, req_type, selection, factor=1, defry=0, prompt=''):
    if selection == "scale":
        width = int(width * factor)
        height = int(height * factor)
    image = image.resize((width, height))
    req_type = {"移除背景": "bg-removal", "素描": "sketch", "线稿": "lineart", "上色": "colorize", "更改表情": "emotion", "去聊天框": "declutter"}[req_type]
    base64img = image2base64(image)
    payload = {"image": base64img, "width": width, "height": height, "req_type": req_type}
    if req_type == "colorize" or req_type == "emotion":
        payload["defry"] = defry
        payload["prompt"] = prompt
    
    try:
        response = requests.post("https://image.novelai.net/ai/augment-image", json=payload, headers=headers, timeout=60)
    except:
        raise gr.Error('NAI response timeout')

    # Process the response
    if response.headers.get('Content-Type') == 'binary/octet-stream':
        zipfile_in_memory = io.BytesIO(response.content)
        with zipfile.ZipFile(zipfile_in_memory, 'r') as zip_ref:
            if len(zip_ref.namelist()):
                images = []
                for file_name in zip_ref.namelist():
                    with zip_ref.open(file_name) as file:
                        images.append(image_from_bytes(file.read()))
                return images
            else:
                messages = json.loads(response.content)
                raise gr.Error(str(messages["statusCode"]) + ": " + messages["message"])
    else:
        messages = json.loads(response.content)
        raise gr.Error(str(messages["statusCode"]) + ": " + messages["message"])
    
def vibe_encode(model, images, extracts):
    b64 = [image2base64(i[0]) for i in images]
    vibes = []
    for i, e in zip(b64, extracts):
        try:
            response = requests.post("https://image.novelai.net/ai/encode-vibe", json={'image': i, 'information_extracted': e, 'model': model}, headers=headers, timeout=60)
        except:
            raise gr.Error('NAI response timeout')
        if response.headers.get('Content-Type') == 'application/binary':
            vibe = base64.b64encode(response.content).decode()
            vibes.append(vibe)
        else:
            messages = json.loads(response.content)
            raise gr.Error(str(messages["statusCode"]) + ": " + messages["message"])
    return b64, vibes

def vibe_to_json(model, ref_image, info_extract, b64img, b64enc):
    data = {"identifier": "novelai-vibe-transfer", "version": 1, "type": "image", "image": b64img}
    data['id'] = hashlib.sha256(b64img.encode()).hexdigest()
    data['encodings'] = {}
    data['encodings']['v4full' if model.endswith('full') else 'v4curated'] = {hashlib.sha256(b64enc.encode()).hexdigest(): {'encoding': b64enc, 'params': {'information_extracted': info_extract}}}
    data['name'] = data['id'][:6] + '-' + data['id'][-6:]
    w, h  = ref_image.size
    scale = 256 / max(w, h)
    resized_w = math.floor(w * scale)
    resized_h = math.floor(h * scale)
    thumbnail = image2base64(ref_image.resize((resized_w, resized_h)))
    data['thumbnail'] = 'data:image/jpeg;base64,' + thumbnail
    data['createdAt'] = round(time.time())
    data['importInfo'] = {'model': model, 'infomation_extracted': info_extract, 'strength': 0.6}
    return json.dumps(data), data['name']