User history v2

.gitignore

@@ -137,4 +137,8 @@ dmypy.json
 .ruff_cache
 
 # Spell checker config
-cspell.json
+cspell.json
+
+mock
+user_history
+_history_snapshots

app.py

@@ -7,26 +7,31 @@ import tempfile
 import gradio as gr
 from gradio_client import Client
 
+import user_history
+
 
 client = Client("runwayml/stable-diffusion-v1-5")
 
 
-def generate(prompt: str) -> tuple[str, list[str]]:
-    negative_prompt = ""
-    guidance_scale = 9.0
+def generate(prompt: str, profile: gr.OAuthProfile | None) -> tuple[str, list[str]]:
     out_dir = client.predict(prompt, fn_index=1)
 
-    config = {
+    metadata = {
         "prompt": prompt,
-        "negative_prompt": negative_prompt,
-        "guidance_scale": guidance_scale,
+        "negative_prompt": "",
+        "guidance_scale": 0.9,
     }
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as config_file:
-        json.dump(config, config_file)
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as metadata_file:
+        json.dump(metadata, metadata_file)
 
     with (pathlib.Path(out_dir) / "captions.json").open() as f:
         paths = list(json.load(f).keys())
-    return paths
+
+    # Saving user history
+    for path in paths:
+        user_history.save_image(label=prompt, image=path, profile=profile, metadata=metadata)
+
+    return paths  # type: ignore
 
 
 with gr.Blocks(css="style.css") as demo:
@@ -39,15 +44,13 @@ with gr.Blocks(css="style.css") as demo:
             height="600px",
             object_fit="scale-down",
         )
+    prompt.submit(fn=generate, inputs=prompt, outputs=gallery)
 
-    prompt.submit(
-        fn=generate,
-        inputs=prompt,
-        outputs=gallery,
-    )
-
+with gr.Blocks() as demo_with_history:
+    with gr.Tab("App"):
+        demo.render()
     with gr.Tab("Past generations"):
-        gr.Markdown("building...")
+        user_history.render()
 
 if __name__ == "__main__":
-    demo.launch()
+    demo_with_history.queue().launch()

gallery_history.py

@@ -1,122 +0,0 @@
-"""
-How to use:
-1. Create a Space with a Persistent Storage attached. Filesystem will be available under `/data`.
-2. Add `hf_oauth: true` to the Space metadata (README.md). Make sure to have Gradio>=3.41.0 configured.
-3. Add `HISTORY_FOLDER` as a Space variable (example. `"/data/history"`).
-4. Add `filelock` as dependency in `requirements.txt`.
-5. Add history gallery to your Gradio app:
-    a. Add imports: `from gallery_history import fetch_gallery_history, show_gallery_history`
-    a. Add `history = show_gallery_history()` within `gr.Blocks` context.
-    b. Add `.then(fn=fetch_gallery_history, inputs=[prompt, result], outputs=history)` on the generate event.
-"""
-import json
-import os
-import shutil
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple
-from uuid import uuid4
-
-import gradio as gr
-from filelock import FileLock
-
-
-_folder = os.environ.get("HISTORY_FOLDER")
-if _folder is None:
-    print(
-        "'HISTORY_FOLDER' environment variable not set. User history will be saved "
-        "locally and will be lost when the Space instance is restarted."
-    )
-    _folder = Path(__file__).parent / "history"
-HISTORY_FOLDER_PATH = Path(_folder)
-
-IMAGES_FOLDER_PATH = HISTORY_FOLDER_PATH / "images"
-IMAGES_FOLDER_PATH.mkdir(parents=True, exist_ok=True)
-
-
-def show_gallery_history():
-    gr.Markdown(
-        "## Your past generations\n\n(Log in to keep a gallery of your previous generations."
-        " Your history will be saved and available on your next visit.)"
-    )
-    with gr.Column():
-        with gr.Row():
-            gr.LoginButton(min_width=250)
-            gr.LogoutButton(min_width=250)
-        gallery = gr.Gallery(
-            label="Past images",
-            show_label=True,
-            elem_id="gallery",
-            object_fit="contain",
-            columns=3,
-            height=300,
-            preview=False,
-            show_share_button=False,
-            show_download_button=False,
-        )
-        gr.Markdown("Make sure to save your images from time to time, this gallery may be deleted in the future.")
-        gallery.attach_load_event(fetch_gallery_history, every=None)
-        return gallery
-
-
-def fetch_gallery_history(
-    prompt: Optional[str] = None,
-    result: Optional[Dict] = None,
-    user: Optional[gr.OAuthProfile] = None,
-):
-    if user is None:
-        return []
-    try:
-        if prompt is not None and result is not None:  # None values means no new images
-            return _update_user_history(user["preferred_username"], [(item["name"], prompt) for item in result])
-        else:
-            return _read_user_history(user["preferred_username"])
-    except Exception as e:
-        raise gr.Error(f"Error while fetching history: {e}") from e
-
-
-####################
-# Internal helpers #
-####################
-
-
-def _read_user_history(username: str) -> List[Tuple[str, str]]:
-    """Return saved history for that user."""
-    with _user_lock(username):
-        path = _user_history_path(username)
-        if path.exists():
-            return json.loads(path.read_text())
-        return []  # No history yet
-
-
-def _update_user_history(username: str, new_images: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
-    """Update history for that user and return it."""
-    with _user_lock(username):
-        # Read existing
-        path = _user_history_path(username)
-        if path.exists():
-            images = json.loads(path.read_text())
-        else:
-            images = []  # No history yet
-
-        # Copy images to persistent folder
-        images = [(_copy_image(src_path), prompt) for src_path, prompt in new_images] + images
-
-        # Save and return
-        path.write_text(json.dumps(images))
-        return images
-
-
-def _user_history_path(username: str) -> Path:
-    return HISTORY_FOLDER_PATH / f"{username}.json"
-
-
-def _user_lock(username: str) -> FileLock:
-    """Ensure history is not corrupted if concurrent calls."""
-    return FileLock(f"{_user_history_path(username)}.lock")
-
-
-def _copy_image(src: str) -> str:
-    """Copy image to the persistent storage."""
-    dst = IMAGES_FOLDER_PATH / f"{uuid4().hex}_{Path(src).name}"  # keep file ext
-    shutil.copyfile(src, dst)
-    return str(dst)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.black]
 line-length = 119
-target_version = ['py37', 'py38', 'py39', 'py310']
+target_version = ['py38', 'py39', 'py310']
 preview = true
 
 [tool.mypy]

requirements.txt

@@ -1,5 +1,6 @@
-gradio>=3.44
+gradio[oauth]>=3.44
 huggingface_hub>=0.17
+Pillow
 
 # dev-deps
 ruff

user_history.py

@@ -0,0 +1,286 @@
+import json
+import os
+import shutil
+import warnings
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Tuple
+from uuid import uuid4
+
+import gradio as gr
+import numpy as np
+from filelock import FileLock
+from PIL.Image import Image
+
+
+def setup(
+    folder_path: str | Path | None = None,
+    delete_button: bool = True,
+    export_button: bool = True,
+) -> None:
+    user_history = _UserHistory()
+    user_history.folder_path = _resolve_folder_path(folder_path)
+    user_history.delete_button = delete_button
+    user_history.export_button = export_button
+    user_history.initialized = True
+
+
+def render() -> None:
+    user_history = _UserHistory()
+
+    # initialize with default config
+    if not user_history.initialized:
+        print("Initializing user history with default config. Use `user_history.setup(...)` to customize.")
+        setup()
+
+    # deactivate if no persistent storage
+    if user_history.folder_path is None:
+        gr.Markdown(
+            "User history is deactivated as no Persistent Storage volume has been found. Please contact the Space"
+            " owner to either assign a [Persistent Storage](https://huggingface.co/docs/hub/spaces-storage) or set"
+            " `folder_path` to a temporary folder."
+        )
+        return
+
+    # Render user history tab
+    gr.Markdown(
+        "## Your past generations\n\n(Log in to keep a gallery of your previous generations."
+        " Your history will be saved and available on your next visit.)"
+    )
+    with gr.Row():
+        gr.LoginButton(min_width=250)
+        gr.LogoutButton(min_width=250)
+        refresh_button = gr.Button("Refresh", icon="./assets/icon_refresh.png")
+        export_button = gr.Button("Export", icon="./assets/icon_download.png")
+        delete_button = gr.Button("Delete history", icon="./assets/icon_delete.png")
+
+    # "Export zip" row (hidden by default)
+    with gr.Row():
+        export_file = gr.File(file_count="single", file_types=[".zip"], label="Exported history", visible=False)
+
+    # "Config deletion" row (hidden by default)
+    with gr.Row():
+        confirm_button = gr.Button("Confirm delete all history", variant="stop", visible=False)
+        cancel_button = gr.Button("Cancel", visible=False)
+
+    # Gallery
+    gallery = gr.Gallery(
+        label="Past images",
+        show_label=True,
+        elem_id="gallery",
+        object_fit="contain",
+        columns=5,
+        height=600,
+        preview=False,
+        show_share_button=False,
+        show_download_button=False,
+    )
+    gr.Markdown("Make sure to save your images from time to time, this gallery may be deleted in the future.")
+    gallery.attach_load_event(_fetch_user_history, every=None)
+
+    # Interactions
+    refresh_button.click(fn=_fetch_user_history, inputs=[], outputs=[gallery], queue=False)
+    export_button.click(fn=_export_user_history, inputs=[], outputs=[export_file], queue=False)
+
+    # Taken from https://github.com/gradio-app/gradio/issues/3324#issuecomment-1446382045
+    delete_button.click(
+        lambda: [gr.update(visible=True), gr.update(visible=True)],
+        outputs=[confirm_button, cancel_button],
+        queue=False,
+    )
+    cancel_button.click(
+        lambda: [gr.update(visible=False), gr.update(visible=False)],
+        outputs=[confirm_button, cancel_button],
+        queue=False,
+    )
+    confirm_button.click(_delete_user_history).then(
+        lambda: [gr.update(visible=False), gr.update(visible=False)],
+        outputs=[confirm_button, cancel_button],
+        queue=False,
+    )
+
+
+def save_image(
+    profile: gr.OAuthProfile | None,
+    image: Image | np.ndarray | str | Path,
+    label: str | None = None,
+    metadata: Dict | None = None,
+):
+    # Ignore images from logged out users
+    if profile is None:
+        return
+    username = profile["preferred_username"]
+
+    # Ignore images if user history not used
+    user_history = _UserHistory()
+    if not user_history.initialized:
+        warnings.warn(
+            "User history is not set in Gradio demo. Saving image is ignored. You must use `user_history.render(...)`"
+            " first."
+        )
+        return
+
+    # Copy image to storage
+    image_path = _copy_image(image, dst_folder=user_history._user_images_path(username))
+
+    # Save new image + metadata
+    if metadata is None:
+        metadata = {}
+    if "datetime" not in metadata:
+        metadata["datetime"] = str(datetime.now())
+    data = {"path": str(image_path), "label": label, "metadata": metadata}
+    with user_history._user_lock(username):
+        with user_history._user_jsonl_path(username).open("a") as f:
+            f.write(json.dumps(data) + "\n")
+
+
+#############
+# Internals #
+#############
+
+
+class _UserHistory(object):
+    _instance = None
+    initialized: bool = False
+
+    folder_path: Path | None
+    delete_button: bool
+    export_button: bool
+
+    def __new__(cls):
+        # Using singleton pattern => we don't want to expose an object (more complex to use) but still want to keep
+        # state between `render` and `save_image` calls.
+        if cls._instance is None:
+            cls._instance = super(_UserHistory, cls).__new__(cls)
+        return cls._instance
+
+    def _user_path(self, username: str) -> Path:
+        if self.folder_path is None:
+            raise Exception("User history is deactivated.")
+        path = self.folder_path / username
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    def _user_lock(self, username: str) -> FileLock:
+        """Ensure history is not corrupted if concurrent calls."""
+        if self.folder_path is None:
+            raise Exception("User history is deactivated.")
+        return FileLock(self.folder_path / f"{username}.lock")  # lock outside of folder => better when exporting ZIP
+
+    def _user_jsonl_path(self, username: str) -> Path:
+        return self._user_path(username) / "history.jsonl"
+
+    def _user_images_path(self, username: str) -> Path:
+        path = self._user_path(username) / "images"
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+
+
+def _fetch_user_history(profile: gr.OAuthProfile | None) -> List[Tuple[str, str]]:
+    """Return saved history for that user, if it exists."""
+    # Cannot load history for logged out users
+    if profile is None:
+        return []
+    username = profile["preferred_username"]
+
+    user_history = _UserHistory()
+    if not user_history.initialized:
+        warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
+        return []
+
+    with user_history._user_lock(username):
+        # No file => no history saved yet
+        jsonl_path = user_history._user_jsonl_path(username)
+        if not jsonl_path.is_file():
+            return []
+
+        # Read history
+        images = []
+        for line in jsonl_path.read_text().splitlines():
+            data = json.loads(line)
+            images.append((data["path"], data["label"] or ""))
+        return list(reversed(images))
+
+
+def _export_user_history(profile: gr.OAuthProfile | None) -> Dict | None:
+    """Zip all history for that user, if it exists and return it as a downloadable file."""
+    # Cannot load history for logged out users
+    if profile is None:
+        return None
+    username = profile["preferred_username"]
+
+    user_history = _UserHistory()
+    if not user_history.initialized:
+        warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
+        return None
+
+    # Zip history
+    with user_history._user_lock(username):
+        path = shutil.make_archive(
+            str(_archives_path() / f"history_{username}"), "zip", user_history._user_path(username)
+        )
+
+    return gr.update(visible=True, value=path)
+
+
+def _delete_user_history(profile: gr.OAuthProfile | None) -> None:
+    """Delete all history for that user."""
+    # Cannot load history for logged out users
+    if profile is None:
+        return
+    username = profile["preferred_username"]
+
+    user_history = _UserHistory()
+    if not user_history.initialized:
+        warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
+        return
+
+    with user_history._user_lock(username):
+        shutil.rmtree(user_history._user_path(username))
+
+
+####################
+# Internal helpers #
+####################
+
+
+def _copy_image(image: Image | np.ndarray | str | Path, dst_folder: Path) -> Path:
+    """Copy image to the images folder."""
+    # Already a path => copy it
+    if isinstance(image, str):
+        image = Path(image)
+    if isinstance(image, Path):
+        dst = dst_folder / f"{uuid4().hex}_{Path(image).name}"  # keep file ext
+        shutil.copyfile(image, dst)
+        return dst
+
+    # Still a Python object => serialize it
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    if isinstance(image, Image):
+        dst = dst_folder / f"{uuid4().hex}.png"
+        image.save(dst)
+        return dst
+
+    raise ValueError(f"Unsupported image type: {type(image)}")
+
+
+def _resolve_folder_path(folder_path: str | Path | None) -> Path | None:
+    if folder_path is not None:
+        return Path(folder_path).expanduser().resolve()
+
+    if os.getenv("SYSTEM") == "spaces":
+        if os.path.exists("/data"):  # Persistent storage is enabled!
+            return Path("/data") / "user_history"
+        else:
+            return None  # No persistent storage => no user history
+
+    # Not in a Space => local folder
+    return Path(__file__).parent / "user_history"
+
+
+def _archives_path() -> Path:
+    # Doesn't have to be on persistent storage as it's only used for download
+    path = Path(__file__).parent / "_history_snapshots"
+    path.mkdir(parents=True, exist_ok=True)
+    return path