multimodalart HF staff commited on
Commit
5d5545c
1 Parent(s): c9217ef

Create user_history.py

Browse files
Files changed (1) hide show
  1. user_history.py +423 -0
user_history.py ADDED
@@ -0,0 +1,423 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ User History is a plugin that you can add to your Spaces to cache generated images for your users.
3
+
4
+ Key features:
5
+ - 🤗 Sign in with Hugging Face
6
+ - Save generated images with their metadata: prompts, timestamp, hyper-parameters, etc.
7
+ - Export your history as zip.
8
+ - Delete your history to respect privacy.
9
+ - Compatible with Persistent Storage for long-term storage.
10
+ - Admin panel to check configuration and disk usage .
11
+
12
+ Useful links:
13
+ - Demo: https://huggingface.co/spaces/Wauplin/gradio-user-history
14
+ - README: https://huggingface.co/spaces/Wauplin/gradio-user-history/blob/main/README.md
15
+ - Source file: https://huggingface.co/spaces/Wauplin/gradio-user-history/blob/main/user_history.py
16
+ - Discussions: https://huggingface.co/spaces/Wauplin/gradio-user-history/discussions
17
+ """
18
+ import json
19
+ import os
20
+ import shutil
21
+ import warnings
22
+ from datetime import datetime
23
+ from functools import cache
24
+ from pathlib import Path
25
+ from typing import Callable, Dict, List, Tuple
26
+ from uuid import uuid4
27
+
28
+ import gradio as gr
29
+ import numpy as np
30
+ import requests
31
+ from filelock import FileLock
32
+ from PIL.Image import Image
33
+
34
+
35
+ def setup(folder_path: str | Path | None = None) -> None:
36
+ user_history = _UserHistory()
37
+ user_history.folder_path = _resolve_folder_path(folder_path)
38
+ user_history.initialized = True
39
+
40
+
41
+ def render() -> None:
42
+ user_history = _UserHistory()
43
+
44
+ # initialize with default config
45
+ if not user_history.initialized:
46
+ print("Initializing user history with default config. Use `user_history.setup(...)` to customize folder_path.")
47
+ setup()
48
+
49
+ # Render user history tab
50
+ gr.Markdown(
51
+ "## Your past generations\n\nLog in to keep a gallery of your previous generations. Your history will be saved"
52
+ " and available on your next visit. Make sure to export your images from time to time as this gallery may be"
53
+ " deleted in the future."
54
+ )
55
+
56
+ if os.getenv("SYSTEM") == "spaces" and not os.path.exists("/data"):
57
+ gr.Markdown(
58
+ "**⚠️ Persistent storage is disabled, meaning your history will be lost if the Space gets restarted."
59
+ " Only the Space owner can setup a Persistent Storage. If you are not the Space owner, consider"
60
+ " duplicating this Space to set your own storage.⚠️**"
61
+ )
62
+
63
+ with gr.Row():
64
+ gr.LoginButton(min_width=250)
65
+ gr.LogoutButton(min_width=250)
66
+ refresh_button = gr.Button(
67
+ "Refresh",
68
+ icon="https://huggingface.co/spaces/Wauplin/gradio-user-history/resolve/main/assets/icon_refresh.png",
69
+ )
70
+ export_button = gr.Button(
71
+ "Export",
72
+ icon="https://huggingface.co/spaces/Wauplin/gradio-user-history/resolve/main/assets/icon_download.png",
73
+ )
74
+ delete_button = gr.Button(
75
+ "Delete history",
76
+ icon="https://huggingface.co/spaces/Wauplin/gradio-user-history/resolve/main/assets/icon_delete.png",
77
+ )
78
+
79
+ # "Export zip" row (hidden by default)
80
+ with gr.Row():
81
+ export_file = gr.File(file_count="single", file_types=[".zip"], label="Exported history", visible=False)
82
+
83
+ # "Config deletion" row (hidden by default)
84
+ with gr.Row():
85
+ confirm_button = gr.Button("Confirm delete all history", variant="stop", visible=False)
86
+ cancel_button = gr.Button("Cancel", visible=False)
87
+
88
+ # Gallery
89
+ gallery = gr.Gallery(
90
+ label="Past images",
91
+ show_label=True,
92
+ elem_id="gallery",
93
+ object_fit="contain",
94
+ columns=5,
95
+ height=600,
96
+ preview=False,
97
+ show_share_button=False,
98
+ show_download_button=False,
99
+ )
100
+ gr.Markdown(
101
+ "User history is powered by"
102
+ " [Wauplin/gradio-user-history](https://huggingface.co/spaces/Wauplin/gradio-user-history). Integrate it to"
103
+ " your own Space in just a few lines of code!"
104
+ )
105
+ gallery.attach_load_event(_fetch_user_history, every=None)
106
+
107
+ # Interactions
108
+ refresh_button.click(fn=_fetch_user_history, inputs=[], outputs=[gallery], queue=False)
109
+ export_button.click(fn=_export_user_history, inputs=[], outputs=[export_file], queue=False)
110
+
111
+ # Taken from https://github.com/gradio-app/gradio/issues/3324#issuecomment-1446382045
112
+ delete_button.click(
113
+ lambda: [gr.update(visible=True), gr.update(visible=True)],
114
+ outputs=[confirm_button, cancel_button],
115
+ queue=False,
116
+ )
117
+ cancel_button.click(
118
+ lambda: [gr.update(visible=False), gr.update(visible=False)],
119
+ outputs=[confirm_button, cancel_button],
120
+ queue=False,
121
+ )
122
+ confirm_button.click(_delete_user_history).then(
123
+ lambda: [gr.update(visible=False), gr.update(visible=False)],
124
+ outputs=[confirm_button, cancel_button],
125
+ queue=False,
126
+ )
127
+
128
+ # Admin section (only shown locally or when logged in as Space owner)
129
+ _admin_section()
130
+
131
+
132
+ def save_image(
133
+ profile: gr.OAuthProfile | None,
134
+ image: Image | np.ndarray | str | Path,
135
+ label: str | None = None,
136
+ metadata: Dict | None = None,
137
+ ):
138
+ # Ignore images from logged out users
139
+ if profile is None:
140
+ return
141
+ username = profile["preferred_username"]
142
+
143
+ # Ignore images if user history not used
144
+ user_history = _UserHistory()
145
+ if not user_history.initialized:
146
+ warnings.warn(
147
+ "User history is not set in Gradio demo. Saving image is ignored. You must use `user_history.render(...)`"
148
+ " first."
149
+ )
150
+ return
151
+
152
+ # Copy image to storage
153
+ image_path = _copy_image(image, dst_folder=user_history._user_images_path(username))
154
+
155
+ # Save new image + metadata
156
+ if metadata is None:
157
+ metadata = {}
158
+ if "datetime" not in metadata:
159
+ metadata["datetime"] = str(datetime.now())
160
+ data = {"path": str(image_path), "label": label, "metadata": metadata}
161
+ with user_history._user_lock(username):
162
+ with user_history._user_jsonl_path(username).open("a") as f:
163
+ f.write(json.dumps(data) + "\n")
164
+
165
+
166
+ #############
167
+ # Internals #
168
+ #############
169
+
170
+
171
+ class _UserHistory(object):
172
+ _instance = None
173
+ initialized: bool = False
174
+ folder_path: Path
175
+
176
+ def __new__(cls):
177
+ # Using singleton pattern => we don't want to expose an object (more complex to use) but still want to keep
178
+ # state between `render` and `save_image` calls.
179
+ if cls._instance is None:
180
+ cls._instance = super(_UserHistory, cls).__new__(cls)
181
+ return cls._instance
182
+
183
+ def _user_path(self, username: str) -> Path:
184
+ path = self.folder_path / username
185
+ path.mkdir(parents=True, exist_ok=True)
186
+ return path
187
+
188
+ def _user_lock(self, username: str) -> FileLock:
189
+ """Ensure history is not corrupted if concurrent calls."""
190
+ return FileLock(self.folder_path / f"{username}.lock") # lock outside of folder => better when exporting ZIP
191
+
192
+ def _user_jsonl_path(self, username: str) -> Path:
193
+ return self._user_path(username) / "history.jsonl"
194
+
195
+ def _user_images_path(self, username: str) -> Path:
196
+ path = self._user_path(username) / "images"
197
+ path.mkdir(parents=True, exist_ok=True)
198
+ return path
199
+
200
+
201
+ def _fetch_user_history(profile: gr.OAuthProfile | None) -> List[Tuple[str, str]]:
202
+ """Return saved history for that user, if it exists."""
203
+ # Cannot load history for logged out users
204
+ if profile is None:
205
+ return []
206
+ username = profile["preferred_username"]
207
+
208
+ user_history = _UserHistory()
209
+ if not user_history.initialized:
210
+ warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
211
+ return []
212
+
213
+ with user_history._user_lock(username):
214
+ # No file => no history saved yet
215
+ jsonl_path = user_history._user_jsonl_path(username)
216
+ if not jsonl_path.is_file():
217
+ return []
218
+
219
+ # Read history
220
+ images = []
221
+ for line in jsonl_path.read_text().splitlines():
222
+ data = json.loads(line)
223
+ images.append((data["path"], data["label"] or ""))
224
+ return list(reversed(images))
225
+
226
+
227
+ def _export_user_history(profile: gr.OAuthProfile | None) -> Dict | None:
228
+ """Zip all history for that user, if it exists and return it as a downloadable file."""
229
+ # Cannot load history for logged out users
230
+ if profile is None:
231
+ return None
232
+ username = profile["preferred_username"]
233
+
234
+ user_history = _UserHistory()
235
+ if not user_history.initialized:
236
+ warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
237
+ return None
238
+
239
+ # Zip history
240
+ with user_history._user_lock(username):
241
+ path = shutil.make_archive(
242
+ str(_archives_path() / f"history_{username}"), "zip", user_history._user_path(username)
243
+ )
244
+
245
+ return gr.update(visible=True, value=path)
246
+
247
+
248
+ def _delete_user_history(profile: gr.OAuthProfile | None) -> None:
249
+ """Delete all history for that user."""
250
+ # Cannot load history for logged out users
251
+ if profile is None:
252
+ return
253
+ username = profile["preferred_username"]
254
+
255
+ user_history = _UserHistory()
256
+ if not user_history.initialized:
257
+ warnings.warn("User history is not set in Gradio demo. You must use `user_history.render(...)` first.")
258
+ return
259
+
260
+ with user_history._user_lock(username):
261
+ shutil.rmtree(user_history._user_path(username))
262
+
263
+
264
+ ####################
265
+ # Internal helpers #
266
+ ####################
267
+
268
+
269
+ def _copy_image(image: Image | np.ndarray | str | Path, dst_folder: Path) -> Path:
270
+ """Copy image to the images folder."""
271
+ # Already a path => copy it
272
+ if isinstance(image, str):
273
+ image = Path(image)
274
+ if isinstance(image, Path):
275
+ dst = dst_folder / f"{uuid4().hex}_{Path(image).name}" # keep file ext
276
+ shutil.copyfile(image, dst)
277
+ return dst
278
+
279
+ # Still a Python object => serialize it
280
+ if isinstance(image, np.ndarray):
281
+ image = Image.fromarray(image)
282
+ if isinstance(image, Image):
283
+ dst = dst_folder / f"{uuid4().hex}.png"
284
+ image.save(dst)
285
+ return dst
286
+
287
+ raise ValueError(f"Unsupported image type: {type(image)}")
288
+
289
+
290
+ def _resolve_folder_path(folder_path: str | Path | None) -> Path:
291
+ if folder_path is not None:
292
+ return Path(folder_path).expanduser().resolve()
293
+
294
+ if os.getenv("SYSTEM") == "spaces" and os.path.exists("/data"): # Persistent storage is enabled!
295
+ return Path("/data") / "_user_history"
296
+
297
+ # Not in a Space or Persistent storage not enabled => local folder
298
+ return Path(__file__).parent / "_user_history"
299
+
300
+
301
+ def _archives_path() -> Path:
302
+ # Doesn't have to be on persistent storage as it's only used for download
303
+ path = Path(__file__).parent / "_user_history_exports"
304
+ path.mkdir(parents=True, exist_ok=True)
305
+ return path
306
+
307
+
308
+ #################
309
+ # Admin section #
310
+ #################
311
+
312
+
313
+ def _admin_section() -> None:
314
+ title = gr.Markdown()
315
+ title.attach_load_event(_display_if_admin(), every=None)
316
+
317
+
318
+ def _display_if_admin() -> Callable:
319
+ def _inner(profile: gr.OAuthProfile | None) -> str:
320
+ if profile is None:
321
+ return ""
322
+ if profile["preferred_username"] in _fetch_admins():
323
+ return _admin_content()
324
+ return ""
325
+
326
+ return _inner
327
+
328
+
329
+ def _admin_content() -> str:
330
+ return f"""
331
+ ## Admin section
332
+
333
+ Running on **{os.getenv("SYSTEM", "local")}** (id: {os.getenv("SPACE_ID")}). {_get_msg_is_persistent_storage_enabled()}
334
+
335
+ Admins: {', '.join(_fetch_admins())}
336
+
337
+ {_get_nb_users()} user(s), {_get_nb_images()} image(s)
338
+
339
+ ### Configuration
340
+
341
+ History folder: *{_UserHistory().folder_path}*
342
+
343
+ Exports folder: *{_archives_path()}*
344
+
345
+ ### Disk usage
346
+
347
+ {_disk_space_warning_message()}
348
+ """
349
+
350
+
351
+ def _get_nb_users() -> int:
352
+ user_history = _UserHistory()
353
+ if not user_history.initialized:
354
+ return 0
355
+ if user_history.folder_path is not None and user_history.folder_path.exists():
356
+ return len([path for path in user_history.folder_path.iterdir() if path.is_dir()])
357
+ return 0
358
+
359
+
360
+ def _get_nb_images() -> int:
361
+ user_history = _UserHistory()
362
+ if not user_history.initialized:
363
+ return 0
364
+ if user_history.folder_path is not None and user_history.folder_path.exists():
365
+ return len([path for path in user_history.folder_path.glob("*/images/*")])
366
+ return 0
367
+
368
+
369
+ def _get_msg_is_persistent_storage_enabled() -> str:
370
+ if os.getenv("SYSTEM") == "spaces":
371
+ if os.path.exists("/data"):
372
+ return "Persistent storage is enabled."
373
+ else:
374
+ return (
375
+ "Persistent storage is not enabled. This means that user histories will be deleted when the Space is"
376
+ " restarted. Consider adding a Persistent Storage in your Space settings."
377
+ )
378
+ return ""
379
+
380
+
381
+ def _disk_space_warning_message() -> str:
382
+ user_history = _UserHistory()
383
+ if not user_history.initialized:
384
+ return ""
385
+
386
+ message = ""
387
+ if user_history.folder_path is not None:
388
+ total, used, _ = _get_disk_usage(user_history.folder_path)
389
+ message += f"History folder: **{used / 1e9 :.0f}/{total / 1e9 :.0f}GB** used ({100*used/total :.0f}%)."
390
+
391
+ total, used, _ = _get_disk_usage(_archives_path())
392
+ message += f"\n\nExports folder: **{used / 1e9 :.0f}/{total / 1e9 :.0f}GB** used ({100*used/total :.0f}%)."
393
+
394
+ return f"{message.strip()}"
395
+
396
+
397
+ def _get_disk_usage(path: Path) -> Tuple[int, int, int]:
398
+ for path in [path] + list(path.parents): # first check target_dir, then each parents one by one
399
+ try:
400
+ return shutil.disk_usage(path)
401
+ except OSError: # if doesn't exist or can't read => fail silently and try parent one
402
+ pass
403
+ return 0, 0, 0
404
+
405
+
406
+ @cache
407
+ def _fetch_admins() -> List[str]:
408
+ # Running locally => fake user is admin
409
+ if os.getenv("SYSTEM") != "spaces":
410
+ return ["FakeGradioUser"]
411
+
412
+ # Running in Space but no space_id => ???
413
+ space_id = os.getenv("SPACE_ID")
414
+ if space_id is None:
415
+ return ["Unknown"]
416
+
417
+ # Running in Space => try to fetch organization members
418
+ # Otherwise, it's not an organization => namespace is the user
419
+ namespace = space_id.split("/")[0]
420
+ response = requests.get(f"https://huggingface.co/api/organizations/{namespace}/members")
421
+ if response.status_code == 200:
422
+ return sorted((member["user"] for member in response.json()), key=lambda x: x.lower())
423
+ return [namespace]