Spaces:
Paused
Paused
import random | |
import io | |
import zipfile | |
import requests | |
import json | |
import base64 | |
from PIL import Image | |
jwt_token = '' | |
url = "https://api.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 generate_novelai_image( | |
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, | |
ref_image=None, | |
info_extract=1, | |
ref_str=0.6, | |
inp_img=None, | |
overlay=False, | |
use_inp=False, | |
inp_str=0.7 | |
): | |
# Assign a random seed if seed is -1 | |
if seed == -1: | |
seed = random.randint(0, 2**32 - 1) | |
# Define the payload | |
payload = { | |
"action": "generate", | |
"input": input_text, | |
"model": "nai-diffusion-3", | |
"parameters": { | |
"width": width, | |
"height": height, | |
"scale": scale, | |
"sampler": sampler, | |
"steps": steps, | |
"n_samples": 1, | |
"ucPreset": 0, | |
"add_original_image": overlay, | |
"cfg_rescale": cfg_rescale, | |
"controlnet_strength": 1, | |
"dynamic_thresholding": dyn_threshold, | |
"params_version": 1, | |
"legacy": False, | |
"legacy_v3_extend": False, | |
"negative_prompt": negative_prompt, | |
"noise": 0, | |
"noise_schedule": schedule, | |
"qualityToggle": True, | |
"reference_information_extracted": info_extract, | |
"reference_strength": ref_str, | |
"seed": seed, | |
"sm": smea, | |
"sm_dyn": dyn, | |
"uncond_scale": 1, | |
} | |
} | |
if ref_image is not None: | |
payload['parameters']['reference_image'] = image2base64(ref_image) | |
if use_inp: | |
payload['action'] = "infill" | |
payload['model'] = 'nai-diffusion-3-inpainting' | |
payload['parameters']['mask'] = image2base64(inp_img['layers'][0]) | |
payload['parameters']['image'] = image2base64(inp_img['background']) | |
payload['parameters']['extra_noise_seed'] = seed | |
payload['parameters']['strength'] = inp_str | |
# Send the POST request | |
response = requests.post(url, json=payload, headers=headers) | |
# Process the response | |
if response.headers.get('Content-Type') == 'application/x-zip-compressed': | |
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: | |
return "NAI doesn't return any images", json.loads(response.content) | |
else: | |
return "Generation failed", json.loads(response.content) | |
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 |