fix(metadata): preserve workflow when recipe images convert to webp

tests/utils/test_exif_utils.py

@@ -1,5 +1,8 @@
 import json
 
+import piexif
+from PIL import Image, PngImagePlugin
+
 from py.utils.exif_utils import ExifUtils
 
 
@@ -59,3 +62,82 @@ def test_append_recipe_metadata_includes_checkpoint(monkeypatch, tmp_path):
         "baseModel": "Illustrious",
     }
     assert payload["base_model"] == "Illustrious"
+
+
+def test_optimize_image_preserves_workflow_when_converting_png_to_webp(tmp_path):
+    image_path = tmp_path / "source.png"
+    png_info = PngImagePlugin.PngInfo()
+    png_info.add_text("parameters", "prompt text\nSteps: 20")
+    png_info.add_text("workflow", json.dumps({"nodes": [{"id": 1}]}))
+
+    Image.new("RGB", (64, 32), color="red").save(image_path, pnginfo=png_info)
+
+    optimized_data, extension = ExifUtils.optimize_image(
+        str(image_path),
+        target_width=32,
+        format="webp",
+        quality=85,
+        preserve_metadata=True,
+    )
+
+    optimized_path = tmp_path / f"optimized{extension}"
+    optimized_path.write_bytes(optimized_data)
+
+    exif_dict = piexif.load(str(optimized_path))
+    assert (
+        exif_dict["0th"][piexif.ImageIFD.ImageDescription].decode("utf-8")
+        == 'Workflow:{"nodes": [{"id": 1}]}'
+    )
+    user_comment = exif_dict["Exif"][piexif.ExifIFD.UserComment]
+    assert user_comment.startswith(b"UNICODE\0")
+    assert user_comment[8:].decode("utf-16be") == "prompt text\nSteps: 20"
+
+
+def test_update_image_metadata_preserves_webp_workflow(tmp_path):
+    image_path = tmp_path / "recipe.webp"
+    exif_dict = {
+        "0th": {
+            piexif.ImageIFD.ImageDescription: 'Workflow:{"nodes":[{"id":1}]}',
+        },
+        "Exif": {
+            piexif.ExifIFD.UserComment: b"UNICODE\0"
+            + "prompt text".encode("utf-16be")
+        },
+    }
+    Image.new("RGB", (32, 32), color="blue").save(
+        image_path, format="WEBP", exif=piexif.dump(exif_dict), quality=85
+    )
+
+    ExifUtils.update_image_metadata(
+        str(image_path), 'prompt text\nRecipe metadata: {"title":"recipe"}'
+    )
+
+    updated_exif = piexif.load(str(image_path))
+    assert (
+        updated_exif["0th"][piexif.ImageIFD.ImageDescription].decode("utf-8")
+        == 'Workflow:{"nodes":[{"id":1}]}'
+    )
+    updated_comment = updated_exif["Exif"][piexif.ExifIFD.UserComment]
+    assert (
+        updated_comment[8:].decode("utf-16be")
+        == 'prompt text\nRecipe metadata: {"title":"recipe"}'
+    )
+
+
+def test_update_image_metadata_preserves_png_workflow(tmp_path):
+    image_path = tmp_path / "recipe.png"
+    png_info = PngImagePlugin.PngInfo()
+    png_info.add_text("parameters", "prompt text")
+    png_info.add_text("workflow", '{"nodes":[{"id":1}]}')
+    Image.new("RGB", (32, 32), color="green").save(image_path, pnginfo=png_info)
+
+    ExifUtils.update_image_metadata(
+        str(image_path), 'prompt text\nRecipe metadata: {"title":"recipe"}'
+    )
+
+    with Image.open(image_path) as img:
+        assert img.info["workflow"] == '{"nodes":[{"id":1}]}'
+        assert (
+            img.info["parameters"]
+            == 'prompt text\nRecipe metadata: {"title":"recipe"}'
+        )