From 6fdbd9d94afcb387ab4a94b1b9c535c8622f8035 Mon Sep 17 00:00:00 2001
From: Victor Giers <mail@victorgiers.com>
Date: Thu, 7 May 2026 11:03:33 +0200
Subject: [PATCH] auto-git:  [change] generate_equirect.py

---
 generate_equirect.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/generate_equirect.py b/generate_equirect.py
index 1bccd79..85a048e 100644
--- a/generate_equirect.py
+++ b/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))