[Feat] Add user history (#4)

- Upload folder using huggingface_hub (4d84406f454a4403ed4182cd9c692c3cb93159cc)
- Update README.md (c34c3cfcc38cbee2e7d40af18e8161161af8f97d)
- Update app.py (9a59460dc560154dc715338e6667eb0cc816b90e)
- Update app.py (9731d8247c28097b1b189f82c84b139cbb834724)
- Update app.py (1b7bc92860b69559ff224d307c6a81c5a599a438)
- Update gallery_history.py (9a1020470b0772477893598918634c9932d8592b)
- Update gallery_history.py (bea6b54c4d326c7b17e437d29e7b2a525f1085ad)
- Update gallery_history.py (bafbdeaea363c06c281b45fa02d33fde029f041c)
- Update gallery_history.py (1b24ae22efb7fe5081c7c51d5ff6380b2242712a)
- Update gallery_history.py (3db09b3c948fea6cc710bca8a861783752b333ee)
- Update gallery_history.py (9904ca7141618d2c24811b68856ae7829c5c2cc9)
- Update gallery_history.py (3b961b17162bed65534a333302070f14a8ca97ec)


Co-authored-by: Apolinário from multimodal AI art <[email protected]>
Co-authored-by: Lucain Pouget <[email protected]>

README.md

@@ -4,10 +4,11 @@ emoji: 🌍
 colorFrom: indigo
 colorTo: red
 sdk: gradio
-sdk_version: 3.39.0
+sdk_version: 3.44.2
 app_file: app.py
 pinned: false
 license: mit
+hf_oauth: true
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -9,6 +9,9 @@ from diffusers.utils import numpy_to_pil
 from diffusers import WuerstchenDecoderPipeline, WuerstchenPriorPipeline
 from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
 from previewer.modules import Previewer
+
+from gallery_history import fetch_gallery_history, show_gallery_history
+
 os.environ['TOKENIZERS_PARALLELISM'] = 'false'
 
 DESCRIPTION = "# Würstchen"
@@ -38,7 +41,7 @@ if torch.cuda.is_available():
     if USE_TORCH_COMPILE:
         prior_pipeline.prior = torch.compile(prior_pipeline.prior, mode="reduce-overhead", fullgraph=True)
         decoder_pipeline.decoder = torch.compile(decoder_pipeline.decoder, mode="reduce-overhead", fullgraph=True)
-    
+
     if PREVIEW_IMAGES:
         previewer = Previewer()
         previewer.load_state_dict(torch.load("previewer/text2img_wurstchen_b_v1_previewer_100k.pt")["state_dict"])
@@ -48,6 +51,7 @@ if torch.cuda.is_available():
             output = previewer(latents)
             output = numpy_to_pil(output.clamp(0, 1).permute(0, 2, 3, 1).cpu().numpy())
             return output
+
     else:
         previewer = None
         callback_prior = None
@@ -96,7 +100,7 @@ def generate(
             if isinstance(r, list):
                 yield r
         prior_output = r
-    
+
     decoder_output = decoder_pipeline(
         image_embeddings=prior_output.image_embeddings,
         prompt=prompt,
@@ -209,6 +213,8 @@ with gr.Blocks(css="style.css") as demo:
         cache_examples=CACHE_EXAMPLES,
     )
 
+    history = show_gallery_history()
+
     inputs = [
             prompt,
             negative_prompt,
@@ -234,6 +240,8 @@ with gr.Blocks(css="style.css") as demo:
         inputs=inputs,
         outputs=result,
         api_name="run",
+    ).then(
+        fn=fetch_gallery_history, inputs=[prompt, result], outputs=history
     )
     negative_prompt.submit(
         fn=randomize_seed_fn,
@@ -246,6 +254,8 @@ with gr.Blocks(css="style.css") as demo:
         inputs=inputs,
         outputs=result,
         api_name=False,
+    ).then(
+        fn=fetch_gallery_history, inputs=[prompt, result], outputs=history
     )
     run_button.click(
         fn=randomize_seed_fn,
@@ -258,7 +268,9 @@ with gr.Blocks(css="style.css") as demo:
         inputs=inputs,
         outputs=result,
         api_name=False,
+    ).then(
+        fn=fetch_gallery_history, inputs=[prompt, result], outputs=history
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()

gallery_history.py

@@ -0,0 +1,129 @@
+"""
+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)

requirements.txt

@@ -5,4 +5,5 @@ invisible-watermark==0.2.0
 Pillow==10.0.0
 torch==2.0.1
 transformers==4.32.1
-compel
+compel
+filelock