fffiloni Wauplin HF staff commited on
Commit
a6eccc7
1 Parent(s): 7c7e1d8

Add User History (#1)

Browse files

- [WIP] (b70163cb001f090b5ea3eb45eea33897835b7bf0)
- Update README.md (43c223f4d7e834990f2e164f5556b6742c9d0a87)
- Update app.py (e7155cf24ab00409d06dd059433384ab8e58b49f)


Co-authored-by: Lucain Pouget <Wauplin@users.noreply.huggingface.co>

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