Create app.py

app.py

@@ -0,0 +1,141 @@
+import datetime
+import gradio as gr
+import requests
+import random
+import io
+import zipfile
+from PIL import Image
+import os
+import numpy as np
+import json
+import boto3
+
+
+# Create an S3 client
+s3 = boto3.client('s3')
+
+def save_to_s3(image_data, payload, file_name):
+    # Define the bucket and the path
+    bucket_name = 'dataset-novelai'
+    folder_name = datetime.datetime.now().strftime("%Y-%m-%d")
+    image_key = f'gradio/{folder_name}/{file_name}.webp'
+    payload_key = f'gradio/{folder_name}/{file_name}.json'
+
+    # Save the image
+    image_data.seek(0)  # Go to the start of the BytesIO object
+    s3.upload_fileobj(image_data, bucket_name, image_key, ExtraArgs={'ContentType': 'image/webp'})
+
+    # Save the payload
+    payload_data = io.BytesIO(payload.encode('utf-8'))
+    s3.upload_fileobj(payload_data, bucket_name, payload_key, ExtraArgs={'ContentType': 'application/json'})
+
+
+# Function to handle the NovelAI API request
+def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scale, sampler):
+    jwt_token = os.environ.get('NAI_API_KEY')
+
+    # Check if quality tags are provided and append to input
+    final_input = input_text
+    if quality_tags:
+        final_input += ", " + quality_tags
+
+    # Assign a random seed if seed is -1
+    if seed == -1:
+        seed = random.randint(0, 2**32 - 1)
+
+    # Define the API URL
+    url = "https://api.novelai.net/ai/generate-image"
+
+    # Set the headers
+    headers = {
+        "Authorization": f"Bearer {jwt_token}",
+        "Content-Type": "application/json",
+        "Origin": "https://novelai.net",
+        "Referer": "https://novelai.net/"
+    }
+
+    # Define the payload
+    payload = {
+        "action": "generate",
+        "input": final_input,
+        "model": "nai-diffusion-3",
+        "parameters": {
+            "width": 832,
+            "height": 1216,
+            "scale": scale,
+            "sampler": sampler,
+            "steps": 28,
+            "n_samples": 1,
+            "ucPreset": 0,
+            "add_original_image": False,
+            "cfg_rescale": 0,
+            "controlnet_strength": 1,
+            "dynamic_thresholding": False,
+            "legacy": False,
+            "negative_prompt": negative_prompt,
+            "noise_schedule": "native",
+            "qualityToggle": True,
+            "scale": 5,
+            "seed": seed,
+            "sm": False,
+            "sm_dyn": False,
+            "steps": 28,
+            "ucPreset": 0,
+            "uncond_scale": 1,
+            "width": 832
+        }
+    }
+
+    # 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:
+                    image = Image.open(file)
+
+                    # Prepare to save the image to S3
+                    buffered = io.BytesIO()
+                    image.save(buffered, format="WEBP", quality=98)
+                    file_name = str(int(datetime.datetime.now().timestamp()))
+                    save_to_s3(buffered, json.dumps(payload, indent=4), file_name)
+
+                    return np.array(image), json.dumps(payload, indent=4)
+
+            else:
+                return "No images found in the zip file.", json.dumps(payload, indent=4)
+    else:
+        return "The response is not a zip file.", json.dumps(payload, indent=4)
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_novelai_image,
+    inputs=[
+        gr.Textbox(label="Input Text"),
+        gr.Textbox(label="Quality Tags", value="best quality, amazing quality, very aesthetic, absurdres"),
+        gr.Slider(minimum=-1, maximum=2**32 - 1, step=1, value=-1, label="Seed"),
+        gr.Textbox(label="Negative Prompt", value="nsfw, lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"),
+        gr.Slider(minimum=1, maximum=20, step=1, value=5, label="Scale"),
+        gr.Dropdown(
+            choices=[
+                "k_euler", "k_euler_ancestral", "k_dpmpp_2s_ancestral", 
+                "k_dpmpp_2m", "k_dpmpp_sde", "ddim_v3"
+            ],
+            value="k_euler",
+            label="Sampler"
+        )
+    ],
+    outputs=[
+        "image",
+        gr.Textbox(label="Submitted Payload")
+    ]
+)
+
+try:
+    iface.launch(share=True)
+except RuntimeError: # use in HF spaces
+    iface.launch()