Preserve workflow metadata in archives

app.py

@@ -23,6 +23,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse
 from gradio.routes import App
 from PIL import Image
+from PIL.PngImagePlugin import PngInfo
 from pydantic import BaseModel, Field
 
 
@@ -227,13 +228,16 @@ def generate_core(prompt, width=1152, height=896, steps=16):
         positive = state["encode"].encode(clip=state["clip"], text=f"{prompt},{SUFFIX}")[0]
         negative = state["encode"].encode(clip=state["negative_clip"], text=NEGATIVE)[0]
         latent = state["latent"].generate(width=width, height=height, batch_size=1)[0]
+        seed = random.randrange(2**64)
         samples = state["sample"].sample(
-            model=state["model"], seed=random.randrange(2**64), steps=steps, cfg=4,
+            model=state["model"], seed=seed, steps=steps, cfg=4,
             sampler_name="euler_ancestral", scheduler="karras", positive=positive,
             negative=negative, latent_image=latent, denoise=1,
         )[0]
         image = state["decode"].decode(vae=state["vae"], samples=samples)[0][0]
         result = Image.fromarray(np.clip(image.detach().cpu().numpy() * 255, 0, 255).astype(np.uint8))
+        from workflow_api import image_metadata
+        result.info.update(image_metadata(prompt, width, height, steps, seed))
         backup_pool.submit(archive_image, result)
         return result
 
@@ -249,7 +253,10 @@ def archive_image(image):
             shutil.copy2(path, local)
         mode = "a" if local.is_file() else "w"
         data = BytesIO()
-        image.save(data, format="PNG")
+        pnginfo = PngInfo()
+        for key in ("prompt", "workflow"):
+            pnginfo.add_text(key, image.info[key])
+        image.save(data, format="PNG", pnginfo=pnginfo)
         with py7zr.SevenZipFile(local, mode=mode, password=PASSWORD) as archive:
             archive.set_encrypted_header(True)
             name = f"{len(archive.getnames()) + 1}ST.png"

workflow_api.py

@@ -588,6 +588,33 @@ def build_extra_pnginfo() -> dict[str, Any] | None:
     }
 
 
+def image_metadata(text, width, height, steps, seed):
+    api = build_workflow()
+    api["118"]["inputs"]["string_a"] = text
+    api["27"]["inputs"].update(width=width, height=height)
+    api["5"]["inputs"].update(seed=seed, steps=steps)
+    api["28"]["inputs"].update(model=["10", 0], clip=["10", 1])
+    api.pop("38")
+
+    workflow = build_extra_pnginfo()["workflow"]
+    nodes = {node["id"]: node for node in workflow["nodes"]}
+    nodes[118]["widgets_values"][0] = text
+    nodes[27]["widgets_values"][:2] = [width, height]
+    nodes[5]["widgets_values"][0] = seed
+    nodes[5]["widgets_values"][2] = steps
+    nodes[10]["outputs"][0]["links"] = [27]
+    nodes[10]["outputs"][1]["links"] = [28]
+    workflow["nodes"] = [node for node in workflow["nodes"] if node["id"] != 38]
+    workflow["links"] = [
+        [link[0], 10, *link[2:]] if link[0] in {27, 28} else link
+        for link in workflow["links"] if link[0] not in {29, 30}
+    ]
+    return {
+        "prompt": json.dumps(api, separators=(",", ":")),
+        "workflow": json.dumps(workflow, separators=(",", ":")),
+    }
+
+
 workflow = build_workflow()
 prompt = json.loads(json.dumps(workflow))
 extra_pnginfo = build_extra_pnginfo()