save metadata

app.py

@@ -1,9 +1,11 @@
 import gc
+import json
 import os
 import random
 import re
 import tempfile
 import urllib.request
+import uuid
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterable, Optional
@@ -65,6 +67,7 @@ else:
 
 import gradio as gr
 from PIL import Image, ImageDraw
+from PIL.PngImagePlugin import PngInfo
 
 # ---------------------------------------------------------------------------
 # Constants -- Anima base model
@@ -636,6 +639,132 @@ def _mock_result(image: Image.Image, prompt: str) -> Image.Image:
     return result
 
 
+# ---------------------------------------------------------------------------
+# Generation metadata -- embedded into every output PNG (tEXt chunks) so a
+# past result can be dropped back in and fully reproduced. Returning a
+# filepath (rather than the PIL Image object) from generate() is required:
+# Gradio's Image component passes str/Path returns straight through to disk
+# untouched, but re-encodes a PIL.Image via its own save path, which would
+# silently drop any metadata not explicitly re-attached at save time.
+# ---------------------------------------------------------------------------
+GENERATION_METADATA_KEY = "anima_metadata"
+_OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "anima_outputs")
+os.makedirs(_OUTPUT_DIR, exist_ok=True)
+
+
+def _build_metadata(
+    *, prompt: str, full_prompt: str, negative_prompt: str, positive_prefix: str,
+    diffusion_key: str, width: int, height: int, steps: int, cfg_scale: float,
+    sigma_shift: Optional[float], seed: int,
+    control_enabled: bool, control_key: str, control_strength: float,
+    lora1_key: str, lora1_url: str, lora1_weight: float,
+    lora2_key: str, lora2_url: str, lora2_weight: float,
+    lora3_key: str, lora3_url: str, lora3_weight: float,
+) -> dict:
+    return {
+        "app": "Anima", "metadata_version": 1,
+        "prompt": prompt, "full_prompt": full_prompt,
+        "negative_prompt": negative_prompt, "positive_prefix": positive_prefix,
+        "diffusion_key": diffusion_key,
+        "width": width, "height": height, "steps": steps, "cfg_scale": cfg_scale,
+        "sigma_shift": sigma_shift, "seed": seed,
+        "control_enabled": bool(control_enabled), "control_key": control_key,
+        "control_strength": control_strength,
+        "lora1_key": lora1_key, "lora1_url": lora1_url, "lora1_weight": lora1_weight,
+        "lora2_key": lora2_key, "lora2_url": lora2_url, "lora2_weight": lora2_weight,
+        "lora3_key": lora3_key, "lora3_url": lora3_url, "lora3_weight": lora3_weight,
+    }
+
+
+def _format_human_readable(meta: dict) -> str:
+    """A1111-style block -- purely for humans/other tools eyeballing the PNG info; restore_from_image() reads the JSON blob instead."""
+    lines = [meta.get("full_prompt", "")]
+    if meta.get("negative_prompt"):
+        lines.append(f"Negative prompt: {meta['negative_prompt']}")
+
+    tail = [
+        f"Steps: {meta['steps']}", f"CFG scale: {meta['cfg_scale']}", f"Seed: {meta['seed']}",
+        f"Size: {meta['width']}x{meta['height']}", f"Model: {meta['diffusion_key']}",
+    ]
+    if meta.get("sigma_shift"):
+        tail.append(f"Sigma shift: {meta['sigma_shift']}")
+    if meta.get("control_enabled"):
+        tail.append(f"ControlNet-LLLite: {meta['control_key']} (strength {meta['control_strength']})")
+    lora_bits = [
+        f"{meta[f'lora{i}_url']} ({meta[f'lora{i}_weight']})" for i in (1, 2, 3)
+        if meta.get(f"lora{i}_key") == "Custom URL" and meta.get(f"lora{i}_url")
+    ]
+    if lora_bits:
+        tail.append(f"LoRA: {', '.join(lora_bits)}")
+    lines.append(", ".join(tail))
+    return "\n".join(lines)
+
+
+def _save_with_metadata(image: Image.Image, meta: dict) -> str:
+    pnginfo = PngInfo()
+    pnginfo.add_text(GENERATION_METADATA_KEY, json.dumps(meta, ensure_ascii=False))
+    pnginfo.add_text("parameters", _format_human_readable(meta))
+    out_path = os.path.join(_OUTPUT_DIR, f"anima_{uuid.uuid4().hex}.png")
+    image.save(out_path, format="PNG", pnginfo=pnginfo)
+    return out_path
+
+
+def restore_from_image(file_path):
+    """Read back everything generate() embedded and refill the form. The
+    control image/mask themselves are never embedded (would bloat every
+    output with a second full-size image), so ControlNet has to be
+    re-uploaded manually -- everything else reproduces the exact result."""
+    if not file_path:
+        raise gr.Error("Upload an image first.")
+    try:
+        with Image.open(file_path) as img:
+            raw = img.info.get(GENERATION_METADATA_KEY)
+    except Exception as exc:
+        raise gr.Error(f"Could not read this image: {exc}")
+    if not raw:
+        raise gr.Error("No Anima generation data found in this image's metadata.")
+    meta = json.loads(raw)
+
+    control_enabled = bool(meta.get("control_enabled", False))
+    control_key = meta.get("control_key", CONTROLNET_CHOICES[1])
+    lora1_key = meta.get("lora1_key", "None")
+    lora2_key = meta.get("lora2_key", "None")
+    lora3_key = meta.get("lora3_key", "None")
+
+    status = "**Settings restored.**"
+    if control_enabled:
+        status += " This used ControlNet-LLLite — re-upload the control image manually."
+
+    return (
+        meta.get("prompt", DEFAULT_PROMPT),
+        meta.get("negative_prompt", DEFAULT_NEGATIVE),
+        meta.get("positive_prefix", DEFAULT_POSITIVE_PREFIX.strip()),
+        meta.get("diffusion_key", DEFAULT_DIFFUSION_KEY),
+        meta.get("width", 1408),
+        meta.get("height", 1408),
+        meta.get("steps", 35),
+        meta.get("cfg_scale", 6.0),
+        meta.get("sigma_shift") or 0.0,
+        gr.Number(value=meta.get("seed", -1), interactive=True),
+        False,  # randomize -> off, we're restoring an explicit seed
+        control_enabled,
+        control_key,
+        gr.Column(visible=control_enabled),
+        _controlnet_mask_visible(control_key),
+        meta.get("control_strength", 1.0),
+        lora1_key,
+        gr.Textbox(value=meta.get("lora1_url", ""), visible=(lora1_key == "Custom URL")),
+        meta.get("lora1_weight", 1.0),
+        lora2_key,
+        gr.Textbox(value=meta.get("lora2_url", ""), visible=(lora2_key == "Custom URL")),
+        meta.get("lora2_weight", 1.0),
+        lora3_key,
+        gr.Textbox(value=meta.get("lora3_url", ""), visible=(lora3_key == "Custom URL")),
+        meta.get("lora3_weight", 1.0),
+        status,
+    )
+
+
 @spaces.GPU(duration=60)
 def generate(
     prompt: str,
@@ -664,6 +793,7 @@ def generate(
     lora3_weight: float,
     progress: gr.Progress = gr.Progress(track_tqdm=False),
 ):
+    raw_prompt = (prompt or "").strip() or DEFAULT_PROMPT
     lora_stack = _build_lora_stack(
         lora1_key, lora1_url, lora1_weight,
         lora2_key, lora2_url, lora2_weight,
@@ -680,6 +810,17 @@ def generate(
 
     use_controlnet = bool(control_enabled) and control_key not in (None, "Off") and control_image is not None
 
+    metadata = _build_metadata(
+        prompt=raw_prompt, full_prompt=prompt, negative_prompt=negative_prompt,
+        positive_prefix=positive_prefix, diffusion_key=diffusion_key,
+        width=width, height=height, steps=steps, cfg_scale=cfg_scale,
+        sigma_shift=sigma_shift_value, seed=seed,
+        control_enabled=use_controlnet, control_key=control_key, control_strength=control_strength,
+        lora1_key=lora1_key, lora1_url=lora1_url, lora1_weight=lora1_weight,
+        lora2_key=lora2_key, lora2_url=lora2_url, lora2_weight=lora2_weight,
+        lora3_key=lora3_key, lora3_url=lora3_url, lora3_weight=lora3_weight,
+    )
+
     if DEV_MODE:
         base = control_image if (use_controlnet and control_image is not None) else Image.new(
             "RGB", (width, height), (26, 26, 46)
@@ -690,7 +831,8 @@ def generate(
             f"**[DEV MODE]** seed `{seed}` · `{width}×{height}` · "
             f"ControlNet: `{control_key if use_controlnet else 'off'}` · LoRA: `{lora_summary}`"
         )
-        return result, seed, f"{result.width}×{result.height}", info
+        out_path = _save_with_metadata(result, metadata)
+        return out_path, seed, f"{result.width}×{result.height}", info
 
     gc.collect()
     torch.cuda.empty_cache()
@@ -730,7 +872,8 @@ def generate(
         f"**ControlNet-LLLite:** `{control_key + f' (strength {control_strength})' if use_controlnet else 'off'}`  \n"
         f"**LoRA:** `{lora_summary}`"
     )
-    return image, seed, f"{image.width}×{image.height}", info
+    out_path = _save_with_metadata(image, metadata)
+    return out_path, seed, f"{image.width}×{image.height}", info
 
 
 # ---------------------------------------------------------------------------
@@ -1052,6 +1195,17 @@ with gr.Blocks(title="Anima", theme=_theme, css=CSS) as demo:
                     result_dims = gr.Textbox(label="Resolution", value="—", interactive=False, elem_classes="meta-tile")
                 info_md = gr.Markdown(elem_id="info-box")
 
+            with gr.Column(elem_classes="q-card"):
+                gr.HTML('<div class="card-label">Restore settings</div>')
+                gr.HTML(
+                    '<p class="card-note">Every image Anima generates has its full settings '
+                    '(prompt, seed, size, model, LoRA…) embedded in its PNG data. Drop one back in '
+                    'to restore them — <b>ControlNet control image excepted</b>, re-upload that '
+                    'yourself.</p>'
+                )
+                restore_upload = gr.Image(type="filepath", show_label=False, height=110, elem_id="restore-drop")
+                restore_status = gr.Markdown()
+
     gr.HTML(
         '<footer class="app-footer">Built with Gradio &middot; '
         '<span>circlestone-labs/Anima &middot; kohya-ss/Anima-LLLite</span></footer>'
@@ -1066,6 +1220,21 @@ with gr.Blocks(title="Anima", theme=_theme, css=CSS) as demo:
     lora2_key.change(fn=_slot_url_visible, inputs=[lora2_key], outputs=[lora2_url])
     lora3_key.change(fn=_slot_url_visible, inputs=[lora3_key], outputs=[lora3_url])
 
+    restore_upload.upload(
+        fn=restore_from_image,
+        inputs=[restore_upload],
+        outputs=[
+            prompt, negative_prompt, positive_prefix,
+            diffusion_key,
+            width, height, steps, cfg_scale, sigma_shift, seed, randomize,
+            control_enabled, control_key, controlnet_panel, control_mask, control_strength,
+            lora1_key, lora1_url, lora1_weight,
+            lora2_key, lora2_url, lora2_weight,
+            lora3_key, lora3_url, lora3_weight,
+            restore_status,
+        ],
+    )
+
     generate_btn.click(
         fn=generate,
         inputs=[