Add User History V0.2 and assets

app.py

@@ -423,7 +423,7 @@ def ui(**kwargs):
             This is your private demo for [UnlimitedMusicGen](https://github.com/Oncorporation/audiocraft), a simple and controllable model for music generation
             presented at: ["Simple and Controllable Music Generation"](https://huggingface.co/papers/2306.05284)
             
-            Disclaimer: This won't run on CPU only. Clone this App and run on GPU instance!
+                Disclaimer: This won't run on CPU only. Clone this App and run on GPU instance!
                         
             Todo: Working on improved Interrupt.
             Theme Available at ["Surn/Beeuty"](https://huggingface.co/spaces/Surn/Beeuty)

user_history.py

@@ -0,0 +1,598 @@
+"""
+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 image, video, audio and document files 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
+"""
+
+__version__ = "0.2.0"
+
+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
+import filetype
+import wave
+from mutagen.mp3 import MP3, EasyMP3
+import torchaudio
+import subprocess
+
+
+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
+
+
+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="./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="gradio_user_history_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")
+            
+def save_file(
+    profile: gr.OAuthProfile | None,
+    image: Image | np.ndarray | str | Path | None = None,
+    video: str | Path | None = None,
+    audio: str | Path | None = None,
+    document: str | Path | None = None,
+    label: str | None = None,
+    metadata: Dict | None = None,
+):
+    # Ignore files from logged out users
+    if profile is None:
+        return
+    username = profile["preferred_username"]
+
+    # Ignore files if user history not used
+    user_history = _UserHistory()
+    if not user_history.initialized:
+        warnings.warn(
+            "User history is not set in Gradio demo. Saving files is ignored. You must use `user_history.render(...)`"
+            " first."
+        )
+        return
+    
+    # Save new files + metadata
+    if metadata is None:
+        metadata = {}
+    if "datetime" not in metadata:
+        metadata["datetime"] = str(datetime.now())
+        
+    # Copy image to storage
+    image_path = None
+    if image is not None:
+        image_path = _copy_image(image, dst_folder=user_history._user_images_path(username))
+        image_path = _add_metadata(image_path, metadata)
+
+    # Copy video to storage
+    if video is not None:    
+        video_path = _copy_file(video, dst_folder=user_history._user_file_path(username, "videos"))
+        video_path = _add_metadata(video_path, metadata)
+
+    # Copy audio to storage
+    if audio is not None:     
+        audio_path = _copy_file(audio, dst_folder=user_history._user_file_path(username, "audios"))
+        audio_path = _add_metadata(audio_path, metadata)
+    
+    # Copy document to storage
+    if document is not None:     
+        document_path = _copy_file(document, dst_folder=user_history._user_file_path(username, "documents"))
+        document_path = _add_metadata(document_path, metadata)
+
+    # Save Json file
+    data = {"image_path": str(image_path), "video_path": str(video_path), "audio_path": str(audio_path), "document_path": str(document_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 _user_file_path(self, username: str, filetype: str = "images") -> Path:        
+        path = self._user_path(username) / filetype
+        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:
+    try:
+        """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"Path(file).name}_{uuid4().hex}.png"
+            image.save(dst)
+            return dst
+
+        raise ValueError(f"Unsupported image type: {type(image)}")
+    
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        if not isinstance(dst, Path):
+            dst = Path(image)
+        return dst  # Return the original file_location if an error occurs
+
+def _copy_file(file: any | np.ndarray | str | Path, dst_folder: Path) -> Path:
+    try:
+        """Copy file to the appropriate folder."""
+        # Already a path => copy it
+        if isinstance(file, str):
+            file = Path(file)
+        if isinstance(file, Path):
+            dst = dst_folder / f"{file.stem}_{uuid4().hex}{file.suffix}"  # keep file ext
+            shutil.copyfile(file, dst)
+            return dst
+
+        # Still a Python object => serialize it
+        if isinstance(file, np.ndarray):
+            file = Image.fromarray(file)
+            dst = dst_folder / f"{file.filename}_{uuid4().hex}{file.suffix}"
+            file.save(dst)
+            return dst
+
+        # try other file types
+        kind = filetype.guess(file)
+        if kind is not None:
+            dst = dst_folder / f"{Path(file).stem}_{uuid4().hex}.{kind.extension}"
+            shutil.copyfile(file, dst)
+            return dst
+        raise ValueError(f"Unsupported file type: {type(file)}")
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        if not isinstance(dst, Path):
+            dst = Path(file)
+        return dst  # Return the original file_location if an error occurs
+
+
+def _add_metadata(file_location: Path, metadata: Dict[str, Any]) -> Path:
+    try:
+        file_type = file_location.suffix
+        valid_file_types = [".wav", ".mp3", ".mp4", ".png"]
+        if file_type not in valid_file_types:
+            raise ValueError("Invalid file type. Valid file types are .wav, .mp3, .mp4, .png")
+
+        if file_type == ".wav":
+            # Open and process .wav file
+            with wave.open(file_location, 'rb') as wav_file:
+                # Get the current metadata
+                current_metadata = {key: value for key, value in wav_file.getparams()._asdict().items() if isinstance(value, (int, float))}
+                
+                # Update metadata
+                current_metadata.update(metadata)
+
+                # Reopen the WAV file in write mode
+                with wave.open(file_location, 'wb') as wav_output_file:
+                    # Set the new metadata
+                    wav_output_file.setparams(wav_file.getparams())
+
+            # Save the WAV file (overwriting the previous version)
+            wav_output_file.close()
+        elif file_type == ".mp3":
+            # Open and process .mp3 file
+            audio = EasyMP3(file_location)
+
+            # Add metadata to the file
+            for key, value in metadata.items():
+                audio[key] = value
+
+            # Save the MP3 file (overwriting the previous version)
+            audio.save()
+        elif file_type == ".mp4":
+            # Open and process .mp4 file
+            # Add metadata to the file
+            wav_file_location = file_location.with_suffix(".wav")
+            wave_exists = wav_file_location.exists()
+            if not wave_exists:
+                # Use torchaudio to create the WAV file if it doesn't exist
+                audio, sample_rate = torchaudio.load(file_location, normalize=True)
+                torchaudio.save(wav_file_location, audio, sample_rate, format='wav')
+
+            # Use ffmpeg to add metadata to the video file
+            metadata_args = [f"{key}={value}" for key, value in metadata.items()]
+            ffmpeg_metadata = ":".join(metadata_args)
+            ffmpeg_cmd = f'ffmpeg -i "{file_location}" -i "{wav_file_location}" -map 0:v:0 -map 1:a:0 -c:v copy -c:a aac -metadata "{ffmpeg_metadata}" "{file_location}"'
+            subprocess.run(ffmpeg_cmd, shell=True, check=True)
+
+            # Remove temporary WAV file
+            if not wave_exists:
+                wav_file_location.unlink()
+        elif file_type == ".png":
+            # Open and process .png file
+            image = Image.open(file_location)
+            exif_data = image.info.get("exif", {})
+            exif_data.update(metadata)
+            # Add metadata to the file
+            image.save(file_location, exif=exif_data)
+
+        return file_location  # Return the path to the modified file
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return file_location  # Return the original file_location if an error occurs
+   
+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("_user_history").resolve()
+
+
+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 and user_history.folder_path.exists():
+        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 and user_history.folder_path.exists():
+        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]