auto-git:

[change] generate_equirect.py

generate_equirect.py

@@ -16,6 +16,7 @@ import shutil
 import tempfile
 import torch
 from PIL import Image, ImageDraw
+from PIL.PngImagePlugin import PngInfo
 import numpy as np
 
 from diffusers import (
@@ -62,6 +63,18 @@ def next_filename(output_dir: str, base: str, width: int, height: int) -> str:
         i += 1
 
 
+def save_png_with_prompt(img: Image.Image, out_path: str, prompt: str) -> None:
+    pnginfo = PngInfo()
+    pnginfo.add_text("prompt", prompt)
+
+    save_kwargs = {"pnginfo": pnginfo}
+    for key in ("icc_profile", "exif", "dpi"):
+        if key in img.info:
+            save_kwargs[key] = img.info[key]
+
+    img.save(out_path, **save_kwargs)
+
+
 def shift_image(img: Image.Image, shift: int) -> Image.Image:
     w, h = img.size
     out = Image.new("RGB", (w, h))