File size: 4,057 Bytes
a411e12
 
 
2af6387
39661d4
f5f05d8
a411e12
6c62d51
a411e12
 
76a000e
bec6ead
2af6387
a411e12
 
 
2af6387
a411e12
 
 
2af6387
a411e12
 
 
 
 
 
2af6387
a411e12
 
 
 
 
 
 
 
 
 
 
 
 
c34bc06
 
 
eef13e5
 
cf5a805
 
 
 
a411e12
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eef13e5
a411e12
 
 
f5f05d8
a411e12
f5f05d8
a411e12
eef13e5
a411e12
 
c34bc06
 
a411e12
 
 
 
c34bc06
a411e12
 
c34bc06
 
9181b21
f5f05d8
5aba480
9e64d02
0d3d342
d46ff0f
cf5a805
eef13e5
 
 
 
a411e12
82e8b9e
825d6b5
82e8b9e
e0d733d
a411e12
 
cb7ba0d
a411e12
 
 
 
 
 
 
39661d4
a411e12
39661d4
a411e12
 
 
 
f5f05d8
 
 
 
 
 
 
9e64d02
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import random
import io
import zipfile
import requests
import json
import base64

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 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_images=None,
    info_extracts=[],
    ref_strs=[],
    i2i_image=None,
    i2i_str=0.7,
    i2i_noise=0,
    overlay=True,
    inp_img=None,
    selection='i2i'
):
    # 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": True,
            "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": i2i_noise,
            "noise_schedule": schedule,
            "qualityToggle": True,
            "reference_information_extracted_multiple": info_extracts,
            "reference_strength_multiple": ref_strs,
            "seed": seed,
            "sm": smea,
            "sm_dyn": dyn,
            "uncond_scale": 1,
            "add_original_image": overlay
        }
    }
    if ref_images is not None:
        payload['parameters']['reference_image_multiple'] = [image2base64(image[0]) for image in ref_images]
    if selection == 'inp' and inp_img['background'].getextrema()[3][1] > 0:
        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
    if i2i_image is not None and selection == 'i2i':
        payload['action'] = "img2img"
        payload['parameters']['image'] = image2base64(i2i_image)
        payload['parameters']['strength'] = i2i_str
        payload['parameters']['extra_noise_seed'] = seed
    # Send the POST request
    try:
        response = requests.post(url, json=payload, headers=headers, timeout=180)
    except:
        return None, {'message': '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:
                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