diff --git a/images/example.png b/images/example.png index fdd1c39..32a7b0f 100644 Binary files a/images/example.png and b/images/example.png differ diff --git a/images/screenshot.png b/images/screenshot.png deleted file mode 100644 index d4fd4d6..0000000 Binary files a/images/screenshot.png and /dev/null differ diff --git a/images/test.png b/images/test.png index 2fdaace..ba5c598 100644 Binary files a/images/test.png and b/images/test.png differ diff --git a/images/test_original.png b/images/test_original.png index 5eb9606..779799d 100644 Binary files a/images/test_original.png and b/images/test_original.png differ diff --git a/png_info_editor.pyw b/png_info_editor.pyw index 840f8e5..b53a58b 100644 --- a/png_info_editor.pyw +++ b/png_info_editor.pyw @@ -7,14 +7,18 @@ from PIL import Image, PngImagePlugin START = "1.0" ERROR = "Oops" +params = None + def select_source_file(): - global source_label, textbox + global source_label, textbox, params filename = filedialog.askopenfilename(initialdir=os.getcwd(), title="Select source file", filetypes=(("PNG", "*.png"),)) if filename: with Image.open(filename, formats=("png",)) as source: - info = source.info.get("parameters", None) - if not info: - return messagebox.showerror(title="Oops", message="No PNG info found in this image.") + if source.info.get("Title", None) == "AI generated image": + params = source.info.copy() + info = source.info.get("Description", None) + else: + return messagebox.showerror(title="Oops", message="No NovelAI3 data found in this image.") source_label.config(text=f"{filename.split('/')[-1]}") textbox.delete(START, tk.END) textbox.insert(START, info) @@ -28,20 +32,23 @@ def select_target_file(): target_file = filename def copy_png_info(): - global target_file, textbox + global target_file, textbox, params info = textbox.get(START, tk.END).strip() if not info: - return messagebox.showerror(title=ERROR, message="Please select a file to copy PNG info from, or write your own.") + return messagebox.showerror(title=ERROR, message="Please select a file to copy NovelAI3 data from.") if not target_file: - return messagebox.showerror(title=ERROR, message="Please select a target image to save the PNG info to.") + return messagebox.showerror(title=ERROR, message="Please select a target image to save the NovelAI3 data to.") try: target = Image.open(target_file, formats=("png",)) except: return messagebox.showerror(title=ERROR, message="The target image might've been moved or is invalid.") + if not params: + return messagebox.showerror(title=ERROR, message="You must first select a source image to copy the original NovelAI3 data from.") pnginfo = PngImagePlugin.PngInfo() - pnginfo.add_text("parameters", info) + for key, val in params.items(): + pnginfo.add_text(key, val) target.save(target_file, format="png", pnginfo=pnginfo) - messagebox.showinfo(title="Success", message=f"PNG info copied to {target_file.split('/')[-1]}") + messagebox.showinfo(title="Success", message=f"NovelAI3 data copied to {target_file.split('/')[-1]}") def apply_color(event=None): global textbox, tags @@ -54,14 +61,14 @@ def apply_color(event=None): textbox.tag_add(tag, f"{i+1}.{match.start(1)}", f"{i+1}.{match.end(1)}") tags = { - "blue": r"(?:^|, )([\w ]+):", - "purple": r"(<\w+:[^>]+>)", + "blue": r"({|})", + "dark red": r"(\[|\])", } target_file = "" window = tk.Tk() -window.title("PNG Info editor") +window.title("NovelAI3 prompt editor") boxtop = tk.Frame(window) boxmiddle = tk.Frame(window) @@ -72,7 +79,7 @@ boxbottom.grid(row=2, padx=5, pady=(0, 10)) for i in range(5): boxmiddle.columnconfigure(i, minsize=100) -tk.Label(boxtop, text="This program lets you copy and edit Stable Diffusion Generation data between images.").pack() +tk.Label(boxtop, text="This program lets you copy and edit NovelAI3 prompts between images.").pack() tk.Button(boxmiddle, text="Copy from image", command=select_source_file).grid(row=0, column=1) tk.Button(boxmiddle, text="Select target image", command=select_target_file).grid(row=0, column=3) source_label = tk.Label(boxmiddle, text="No file selected.")