Spaces:
Paused
Paused
| 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_image=None, | |
| info_extract=1, | |
| ref_str=0.6, | |
| 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": info_extract, | |
| "reference_strength": ref_str, | |
| "seed": seed, | |
| "sm": smea, | |
| "sm_dyn": dyn, | |
| "uncond_scale": 1, | |
| "overlay": True | |
| } | |
| } | |
| if ref_image is not None: | |
| payload['parameters']['reference_image'] = image2base64(ref_image) | |
| if selection == 'inp' and inp_img['background'].getextrema()[3][1] > 0: | |
| payload['action'] = "infill" | |
| payload['model'] = 'nai-diffusion-3-inpainting' | |
| payload['parameters']['mask'] = image2base64(convert_mask(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=20) | |
| 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 | |
| def convert_mask(img): | |
| a = img.split()[3] | |
| base = Image.new('RGB', img.size, 'white') | |
| base.putalpha(a) | |
| return base |