server/nai_api.py

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