Spaces:
Running on CPU Upgrade

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