User history V2 (#101)

- [WIP] (1615e088806094bb6947ed7bb99b2d10760dea13)
- Delete gallery_history.py (57d05a070108a1e900aa81aeff163812e3cd69dc)


Co-authored-by: Lucain Pouget <[email protected]>

app.py

@@ -16,7 +16,7 @@ from diffusers import (
 )
 import time
 from share_btn import community_icon_html, loading_icon_html, share_js
-from gallery_history import fetch_gallery_history, show_gallery_history
+import user_history
 from illusion_style import css
 
 BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"
@@ -117,7 +117,8 @@ def inference(
     upscaler_strength: float = 0.5,
     seed: int = -1,
     sampler = "DPM++ Karras SDE",
-    progress = gr.Progress(track_tqdm=True)
+    progress = gr.Progress(track_tqdm=True),
+    profile: gr.OAuthProfile | None = None,
 ):
     start_time = time.time()
     start_time_struct = time.localtime(start_time)
@@ -165,9 +166,28 @@ def inference(
     end_time_struct = time.localtime(end_time)
     end_time_formatted = time.strftime("%H:%M:%S", end_time_struct)
     print(f"Inference ended at {end_time_formatted}, taking {end_time-start_time}s")
+
+    # Save image + metadata
+    user_history.save_image(
+        label=prompt,
+        image=out_image["images"][0],
+        profile=profile,
+        metadata={
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "guidance_scale": guidance_scale,
+            "controlnet_conditioning_scale": controlnet_conditioning_scale,
+            "control_guidance_start": control_guidance_start,
+            "control_guidance_end": control_guidance_end,
+            "upscaler_strength": upscaler_strength,
+            "seed": seed,
+            "sampler": sampler,
+        },
+    )
+
     return out_image["images"][0], gr.update(visible=True), gr.update(visible=True), my_seed
         
-with gr.Blocks(css=css) as app:
+with gr.Blocks() as app:
     gr.Markdown(
         '''
         <center><h1>Illusion Diffusion HQ 🌀</h1></span>  
@@ -205,7 +225,6 @@ with gr.Blocks(css=css) as app:
                 loading_icon = gr.HTML(loading_icon_html)
                 share_button = gr.Button("Share to community", elem_id="share-btn")
 
-    history = show_gallery_history()
     prompt.submit(
         check_inputs,
         inputs=[prompt, control_image],
@@ -226,8 +245,6 @@ with gr.Blocks(css=css) as app:
         outputs=[result_image],
         queue=False,
         postprocess=False
-    ).success(
-        fn=fetch_gallery_history, inputs=[prompt, result_image], outputs=history, queue=False
     )
     run_btn.click(
         check_inputs,
@@ -249,11 +266,16 @@ with gr.Blocks(css=css) as app:
         outputs=[result_image],
         queue=False,
         postprocess=False
-    ).success(
-        fn=fetch_gallery_history, inputs=[prompt, result_image], outputs=history, queue=False
     )
     share_button.click(None, [], [], _js=share_js)
-app.queue(max_size=20)
+
+with gr.Blocks(css=css) as app_with_history:
+    with gr.Tab("Demo"):
+        app.render()
+    with gr.Tab("Past generations"):
+        user_history.render()
+
+app_with_history.queue(max_size=20)
 
 if __name__ == "__main__":
-    app.launch(max_threads=400)
+    app_with_history.launch(max_threads=400)

gallery_history.py

@@ -1,134 +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 numpy as np
-import shutil
-from pathlib import Path
-from PIL import Image
-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"
-if _folder.startswith("/data") and not os.path.exists("/data"):
-    print(
-        f"'HISTORY_FOLDER' environment variable is set to '{_folder}' which doesn't exist. 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=4,
-            height=512,
-            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[np.ndarray] = 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
-            new_image = Image.fromarray(result, 'RGB')
-            return _update_user_history(user["preferred_username"], new_image, prompt)
-        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_image: Image.Image, prompt: 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 image to persistent folder
-        images = [(_copy_image(new_image), prompt)] + 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(new_image: Image.Image) -> str:
-    """Copy image to the persistent storage."""
-    dst = str(IMAGES_FOLDER_PATH / f"{uuid4().hex}.png")
-    new_image.save(dst)
-    return dst

user_history.py

@@ -0,0 +1,524 @@
+"""
+User History is a plugin that you can add to your Spaces to cache generated images for your users.
+
+Key features:
+- 🤗 Sign in with Hugging Face
+- Save generated images with their metadata: prompts, timestamp, hyper-parameters, etc.
+- Export your history as zip.
+- Delete your history to respect privacy.
+- Compatible with Persistent Storage for long-term storage.
+- Admin panel to check configuration and disk usage .
+
+Useful links:
+- Demo: https://huggingface.co/spaces/Wauplin/gradio-user-history
+- README: https://huggingface.co/spaces/Wauplin/gradio-user-history/blob/main/README.md
+- Source file: https://huggingface.co/spaces/Wauplin/gradio-user-history/blob/main/user_history.py
+- Discussions: https://huggingface.co/spaces/Wauplin/gradio-user-history/discussions
+"""
+import json
+import os
+import shutil
+import warnings
+from datetime import datetime
+from functools import cache
+from pathlib import Path
+from typing import Callable, Dict, List, Tuple
+from uuid import uuid4
+
+import gradio as gr
+import numpy as np
+import requests
+from filelock import FileLock
+from PIL.Image import Image
+
+
+def setup(folder_path: str | Path | None = None) -> None:
+    user_history = _UserHistory()
+    user_history.folder_path = _resolve_folder_path(folder_path)
+    user_history.initialized = True
+
+    # TODO: remove this section once all Spaces have migrated
+    _migrate_history()
+
+
+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 folder_path."
+        )
+        setup()
+
+    # Render user history tab
+    gr.Markdown(
+        "## Your past generations\n\nLog in to keep a gallery of your previous generations. Your history will be saved"
+        " and available on your next visit. Make sure to export your images from time to time as this gallery may be"
+        " deleted in the future."
+    )
+
+    if os.getenv("SYSTEM") == "spaces" and not os.path.exists("/data"):
+        gr.Markdown(
+            "**⚠️ Persistent storage is disabled, meaning your history will be lost if the Space gets restarted."
+            " Only the Space owner can setup a Persistent Storage. If you are not the Space owner, consider"
+            " duplicating this Space to set your own storage.⚠️**"
+        )
+
+    with gr.Row():
+        gr.LoginButton(min_width=250)
+        gr.LogoutButton(min_width=250)
+        refresh_button = gr.Button(
+            "Refresh",
+            icon="https://huggingface.co/spaces/Wauplin/gradio-user-history/resolve/main/assets/icon_refresh.png",
+        )
+        export_button = gr.Button(
+            "Export",
+            icon="https://huggingface.co/spaces/Wauplin/gradio-user-history/resolve/main/assets/icon_download.png",
+        )
+        delete_button = gr.Button(
+            "Delete history",
+            icon="https://huggingface.co/spaces/Wauplin/gradio-user-history/resolve/main/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(
+        "User history is powered by"
+        " [Wauplin/gradio-user-history](https://huggingface.co/spaces/Wauplin/gradio-user-history). Integrate it to"
+        " your own Space in just a few lines of code!"
+    )
+    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,
+    )
+
+    # Admin section (only shown locally or when logged in as Space owner)
+    _admin_section()
+
+
+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
+
+    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:
+        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."""
+        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:
+    if folder_path is not None:
+        return Path(folder_path).expanduser().resolve()
+
+    if os.getenv("SYSTEM") == "spaces" and os.path.exists(
+        "/data"
+    ):  # Persistent storage is enabled!
+        return Path("/data") / "_user_history"
+
+    # Not in a Space or Persistent storage not enabled => 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 / "_user_history_exports"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+#################
+# Admin section #
+#################
+
+
+def _admin_section() -> None:
+    title = gr.Markdown()
+    title.attach_load_event(_display_if_admin(), every=None)
+
+
+def _display_if_admin() -> Callable:
+    def _inner(profile: gr.OAuthProfile | None) -> str:
+        if profile is None:
+            return ""
+        if profile["preferred_username"] in _fetch_admins():
+            return _admin_content()
+        return ""
+
+    return _inner
+
+
+def _admin_content() -> str:
+    return f"""
+## Admin section
+
+Running on **{os.getenv("SYSTEM", "local")}** (id: {os.getenv("SPACE_ID")}). {_get_msg_is_persistent_storage_enabled()}
+
+Admins: {', '.join(_fetch_admins())}
+
+{_get_nb_users()} user(s), {_get_nb_images()} image(s)
+
+### Configuration
+
+History folder: *{_UserHistory().folder_path}*
+
+Exports folder: *{_archives_path()}*
+
+### Disk usage
+
+{_disk_space_warning_message()}
+"""
+
+
+def _get_nb_users() -> int:
+    user_history = _UserHistory()
+    if not user_history.initialized:
+        return 0
+    if user_history.folder_path is not None:
+        return len(
+            [path for path in user_history.folder_path.iterdir() if path.is_dir()]
+        )
+    return 0
+
+
+def _get_nb_images() -> int:
+    user_history = _UserHistory()
+    if not user_history.initialized:
+        return 0
+    if user_history.folder_path is not None:
+        return len([path for path in user_history.folder_path.glob("*/images/*")])
+    return 0
+
+
+def _get_msg_is_persistent_storage_enabled() -> str:
+    if os.getenv("SYSTEM") == "spaces":
+        if os.path.exists("/data"):
+            return "Persistent storage is enabled."
+        else:
+            return (
+                "Persistent storage is not enabled. This means that user histories will be deleted when the Space is"
+                " restarted. Consider adding a Persistent Storage in your Space settings."
+            )
+    return ""
+
+
+def _disk_space_warning_message() -> str:
+    user_history = _UserHistory()
+    if not user_history.initialized:
+        return ""
+
+    message = ""
+    if user_history.folder_path is not None:
+        total, used, _ = _get_disk_usage(user_history.folder_path)
+        message += f"History folder: **{used / 1e9 :.0f}/{total / 1e9 :.0f}GB** used ({100*used/total :.0f}%)."
+
+    total, used, _ = _get_disk_usage(_archives_path())
+    message += f"\n\nExports folder: **{used / 1e9 :.0f}/{total / 1e9 :.0f}GB** used ({100*used/total :.0f}%)."
+
+    return f"{message.strip()}"
+
+
+def _get_disk_usage(path: Path) -> Tuple[int, int, int]:
+    for path in [path] + list(
+        path.parents
+    ):  # first check target_dir, then each parents one by one
+        try:
+            return shutil.disk_usage(path)
+        except (
+            OSError
+        ):  # if doesn't exist or can't read => fail silently and try parent one
+            pass
+    return 0, 0, 0
+
+
+@cache
+def _fetch_admins() -> List[str]:
+    # Running locally => fake user is admin
+    if os.getenv("SYSTEM") != "spaces":
+        return ["FakeGradioUser"]
+
+    # Running in Space but no space_id => ???
+    space_id = os.getenv("SPACE_ID")
+    if space_id is None:
+        return ["Unknown"]
+
+    # Running in Space => try to fetch organization members
+    # Otherwise, it's not an organization => namespace is the user
+    namespace = space_id.split("/")[0]
+    response = requests.get(
+        f"https://huggingface.co/api/organizations/{namespace}/members"
+    )
+    if response.status_code == 200:
+        return sorted(
+            (member["user"] for member in response.json()), key=lambda x: x.lower()
+        )
+    return [namespace]
+
+
+################################################################
+# Legacy helpers to migrate image structure to new data format #
+################################################################
+# TODO: remove this section once all Spaces have migrated
+
+
+def _migrate_history():
+    """Script to migrate user history from v0 to v1."""
+    legacy_history_path = _legacy_get_history_folder_path()
+    if not legacy_history_path.exists():
+        return
+
+    error_count = 0
+    for json_path in legacy_history_path.glob("*.json"):
+        username = json_path.stem
+        print(f"Migrating history for user {username}...")
+        error_count += _legacy_move_user_history(username)
+        print("Done.")
+    print(f"Migration complete. {error_count} error(s) happened.")
+
+    if error_count == 0:
+        shutil.rmtree(legacy_history_path, ignore_errors=True)
+
+
+def _legacy_move_user_history(username: str) -> int:
+    history = _legacy_read_user_history(username)
+    error_count = 0
+    for image, prompt in reversed(history):
+        try:
+            save_image(
+                label=prompt, image=image, profile={"preferred_username": username}
+            )
+        except Exception as e:
+            print("Issue while migrating image:", e)
+            error_count += 1
+    return error_count
+
+
+def _legacy_get_history_folder_path() -> Path:
+    _folder = os.environ.get("HISTORY_FOLDER")
+    if _folder is None:
+        _folder = Path(__file__).parent / "history"
+    return Path(_folder)
+
+
+def _legacy_read_user_history(username: str) -> List[Tuple[str, str]]:
+    """Return saved history for that user."""
+    with _legacy_user_lock(username):
+        path = _legacy_user_history_path(username)
+        if path.exists():
+            return json.loads(path.read_text())
+        return []  # No history yet
+
+
+def _legacy_user_history_path(username: str) -> Path:
+    return _legacy_get_history_folder_path() / f"{username}.json"
+
+
+def _legacy_user_lock(username: str) -> FileLock:
+    """Ensure history is not corrupted if concurrent calls."""
+    return FileLock(f"{_legacy_user_history_path(username)}.lock")