Update utils.py

utils.py

@@ -3,6 +3,7 @@ import io
 import zipfile
 import requests
 import json
+import base64
 
 from PIL import Image
 
@@ -38,6 +39,12 @@ def generate_novelai_image(
     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
 ):
     # Assign a random seed if seed is -1
     if seed == -1:
@@ -56,21 +63,32 @@ def generate_novelai_image(
             "steps": steps,
             "n_samples": 1,
             "ucPreset": 0,
-            "add_original_image": False,
+            "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": 1,
+            "reference_strength": 0.6,
             "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-inpaint'
+        payload['parameters']['mask'] = image2base64(inp_img['layers'][0])
+        payload['parameters']['image'] = image2base64(inp_img['background'])
     # Send the POST request
     response = requests.post(url, json=payload, headers=headers)
 
@@ -93,4 +111,11 @@ def generate_novelai_image(
 def image_from_bytes(data):
     img_file = io.BytesIO(data)
     img_file.seek(0)
-    return Image.open(img_file)
+    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