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']