import base64 import io import json import os import random import zipfile import httpx from .utils import extract_png_info def get_api_headers(): api_key = os.environ.get("NAI_API_KEY", "") if not api_key: raise ValueError("NAI_API_KEY environment variable not set") return { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "Origin": "https://novelai.net", "Referer": "https://novelai.net/" } # ────────────────────────────────────────────── # Vibe encoding # ────────────────────────────────────────────── async def encode_vibes(params: dict) -> list: headers = get_api_headers() model = params["model"] images = params["images"] results = [] for i, img_data in enumerate(images): payload = { "image": img_data["b64"], "information_extracted": img_data["info_extract"], "model": model, } async with httpx.AsyncClient(timeout=60.0) as client: response = await client.post( "https://image.novelai.net/ai/encode-vibe", json=payload, headers=headers, ) if response.status_code == 200: # Make sure it isn't a JSON error disguised as 200 try: error_data = json.loads(response.content) if "statusCode" in error_data: raise ValueError( f"{error_data['statusCode']}: " f"{error_data.get('message', 'Unknown error')}" ) except (json.JSONDecodeError, UnicodeDecodeError): pass vibe_b64 = base64.b64encode(response.content).decode() results.append({ "index": i, "b64img": img_data["b64"], "b64enc": vibe_b64, "info_extract": img_data["info_extract"], }) else: try: messages = json.loads(response.content) raise ValueError( f"{messages.get('statusCode', response.status_code)}: " f"{messages.get('message', 'Unknown error')}" ) except json.JSONDecodeError: raise ValueError(f"NAI API error: {response.status_code}") return results # ────────────────────────────────────────────── # Image generation # ────────────────────────────────────────────── async def generate_image(params: dict) -> dict: headers = get_api_headers() model = params.get("model", "nai-diffusion-4-5-full") seed = params.get("seed", -1) if seed == -1: seed = random.randint(0, 4294967295) quality_prompt = params.get("quality_prompt", "") negative_prompt = params.get("negative_prompt", "") full_prompt = params.get("prompt", "") if quality_prompt: full_prompt = f"{full_prompt}, {quality_prompt}" if full_prompt else quality_prompt is_v4 = model.startswith("nai-diffusion-4") # Determine action from extra_input extra_input = params.get("extra_input") action = "generate" if extra_input: mode = extra_input.get("mode", "") if mode == "img2img": action = "img2img" elif mode == "inpaint": action = "infill" payload = { "input": full_prompt, "model": model, "action": action, "parameters": { "width": params.get("width", 1024), "height": params.get("height", 1024), "scale": params.get("cfg_scale", 5.0), "sampler": params.get("sampler", "k_euler"), "steps": params.get("steps", 28), "seed": seed, "n_samples": 1, "negative_prompt": negative_prompt, "ucPreset": 0, "controlnet_strength": 1, "legacy": False, "legacy_uc": False, "legacy_v3_extend": False, "params_version": 3, "qualityToggle": True, "reference_strength_multiple": [], "deliberate_euler_ancestral_bug": False, "prefer_brownian": True, "reference_strength_multiple": [], "normalize_reference_strength_multiple": True, "add_original_image": True, "characterPrompts": [], "use_coords": False, }, } # --- Extra input (img2img / inpaint) --- if extra_input: mode = extra_input.get("mode", "") if mode == "img2img": payload["parameters"]["image"] = extra_input.get("image", "") payload["parameters"]["strength"] = extra_input.get("strength", 0.7) payload["parameters"]["noise"] = extra_input.get("noise", 0) payload["parameters"]["extra_noise_seed"] = seed elif mode == "inpaint": payload['model'] = model.replace('-preview', '') + '-inpainting' payload["parameters"]["image"] = extra_input.get("image", "") payload["parameters"]["mask"] = extra_input.get("mask", "") payload["parameters"]["inpaintImg2ImgStrength"] = extra_input.get("strength", 1) payload["parameters"]["extra_noise_seed"] = seed payload["parameters"]["add_original_image"] = extra_input.get("add_original_image", True) payload["parameters"]["img2img"] = {"strength": extra_input.get("strength", 1), "color_correct": True} # --- Vibe transfer --- vibe = params.get("vibe", {}) vibe_mode = vibe.get("mode", "") if vibe_mode == "v4" and is_v4: vibes_data = vibe.get("vibes", []) if vibes_data: payload["parameters"]["reference_image_multiple"] = [ v["encoding"] for v in vibes_data ] payload["parameters"]["reference_information_extracted_multiple"] = [ v["info_extract"] for v in vibes_data ] payload["parameters"]["reference_strength_multiple"] = [ v["strength"] for v in vibes_data ] elif vibe_mode == "legacy" and not is_v4: images_data = vibe.get("images", []) if images_data: payload["parameters"]["reference_image_multiple"] = [ img["b64"] for img in images_data ] payload["parameters"]["reference_information_extracted_multiple"] = [ img["info_extract"] for img in images_data ] payload["parameters"]["reference_strength_multiple"] = [ img["strength"] for img in images_data ] payload["parameters"]["normalize_reference_strength_multiple"] = vibe.get("normalize_strength", True) # --- Advanced options --- adv = params.get("advanced", {}) noise_schedule = adv.get("scheduler", "karras") sm = adv.get("smea", False) sm_dyn = adv.get("smea_dyn", False) variety = adv.get("variety", False) legacy_v3_extend = adv.get("legacy_v3_extend", False) decrisp = adv.get("decrisp", False) cfg_rescale_val = adv.get("cfg_rescale", 0) payload["parameters"]["noise_schedule"] = noise_schedule payload["parameters"]["sm"] = sm payload["parameters"]["sm_dyn"] = sm_dyn if variety: payload["parameters"]["skip_cfg_above_sigma"] = 58 if model.startswith('nai-diffusion-4-5') else 19 else: payload["parameters"]["skip_cfg_above_sigma"] = None payload["parameters"]["legacy_uc"] = legacy_v3_extend payload["parameters"]["dynamic_thresholding"] = decrisp payload["parameters"]["cfg_rescale"] = cfg_rescale_val # --- Character reference --- reference = params.get("reference") if is_v4 and reference and reference.get("image"): dir_image_b64 = reference["image"] dir_info_extract = reference.get("info_extract", 1) dir_strength = reference.get("strength", 1) dir_style = reference.get("style_mode", "character&style") payload["parameters"]["director_reference_images"] = [dir_image_b64] payload["parameters"]["director_reference_information_extracted"] = [1] payload["parameters"]["director_reference_strength_values"] = [dir_strength] payload["parameters"]["director_reference_descriptions"] = [{ "caption": { "base_caption": dir_style, "char_captions": [], }, "legacy_uc": False, }] payload["parameters"]["director_reference_secondary_strength_values"] = [1 - dir_info_extract] # --- Character prompts --- if is_v4: char_data = params.get("character", {}) char_list = char_data.get("characters", []) use_coords = char_data.get("useCoords", False) character_prompts = [] v4_char_captions = [] v4_neg_char_captions = [] for ch in char_list: cp = { "prompt": ch.get("prompt", ""), "uc": ch.get("uc", ""), "center": ch.get("center", {"x": 0.5, "y": 0.5}), } character_prompts.append(cp) # Build v4 char captions v4_char_captions.append({ "char_caption": ch.get("prompt", ""), "centers": [cp["center"]], }) v4_neg_char_captions.append({ "char_caption": ch.get("uc", ""), "centers": [cp["center"]], }) payload["parameters"].update({ "use_coords": use_coords, "characterPrompts": character_prompts, "v4_prompt": { "caption": { "base_caption": full_prompt, "char_captions": v4_char_captions, }, "use_coords": use_coords, "use_order": True, }, "v4_negative_prompt": { "caption": { "base_caption": negative_prompt, "char_captions": v4_neg_char_captions, } }, }) async with httpx.AsyncClient(timeout=120.0) as client: response = await client.post( "https://image.novelai.net/ai/generate-image", json=payload, headers=headers, ) if response.status_code != 200: error_text = response.text[:500] try: err = json.loads(response.text) error_text = ( f"{err.get('statusCode', response.status_code)}: " f"{err.get('message', error_text)}" ) except Exception: pass raise ValueError(f"NAI API error: {error_text}") zip_bytes = response.content with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf: image_filename = zf.namelist()[0] image_bytes = zf.open(image_filename).read() image_b64 = base64.b64encode(image_bytes).decode("utf-8") info = extract_png_info(image_bytes) info["seed"] = seed return {"image": image_b64, "info": info, "seed": seed} # ────────────────────────────────────────────── # Augment image (Director Tools) # ────────────────────────────────────────────── async def augment_image(params: dict) -> dict: headers = get_api_headers() payload = { "image": params["image"], "width": params["width"], "height": params["height"], "req_type": params["req_type"], } if params["req_type"] in ("colorize", "emotion"): payload["defry"] = params.get("defry", 0) payload["prompt"] = params.get("prompt", "") async with httpx.AsyncClient(timeout=60.0) as client: response = await client.post( "https://image.novelai.net/ai/augment-image", json=payload, headers=headers, ) content_type = response.headers.get("Content-Type", "") if "binary/octet-stream" in content_type or "application/zip" in content_type or response.status_code == 200: try: error_data = json.loads(response.content) if "statusCode" in error_data: raise ValueError( f"{error_data['statusCode']}: {error_data.get('message', 'Unknown error')}" ) except (json.JSONDecodeError, UnicodeDecodeError): pass images_b64 = [] zip_bytes = response.content with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf: for name in zf.namelist(): img_bytes = zf.open(name).read() images_b64.append(base64.b64encode(img_bytes).decode("utf-8")) if not images_b64: raise ValueError("No images in response") return {"images": images_b64} else: try: messages = json.loads(response.content) raise ValueError( f"{messages.get('statusCode', response.status_code)}: " f"{messages.get('message', 'Unknown error')}" ) except json.JSONDecodeError: raise ValueError(f"NAI API error: {response.status_code}")