Add User History

#26
by Wauplin HF staff - opened
Files changed (3) hide show
  1. README.md +1 -0
  2. app.py +23 -4
  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
@@ -14,6 +14,7 @@ import copy
14
  import gc
15
  import pickle
16
  import spaces
 
17
 
18
  lora_list = hf_hub_download(repo_id="multimodalart/LoraTheExplorer", filename="sdxl_loras.json", repo_type="space")
19
 
@@ -93,7 +94,7 @@ div#share-btn-container > div {flex-direction: row;background: black;align-items
93
  original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
94
 
95
  @spaces.GPU
96
- def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lora_2_scale=0.5, seed=-1):
97
 
98
  repo_id_1 = shuffled_items[0]['repo']
99
  repo_id_2 = shuffled_items[1]['repo']
@@ -134,6 +135,18 @@ def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lor
134
  seed = random.randint(0, 2147483647)
135
  generator = torch.Generator(device="cuda").manual_seed(seed)
136
  image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=20, width=768, height=768, generator=generator).images[0]
 
 
 
 
 
 
 
 
 
 
 
 
137
  return image, gr.update(visible=True), seed, gr.update(visible=True, interactive=True), gr.update(visible=False), gr.update(visible=True, interactive=True), gr.update(visible=False)
138
 
139
  def get_description(item):
@@ -176,7 +189,7 @@ def save_preferences(lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, g
176
  def hide_post_gen_info():
177
  return gr.update(visible=False)
178
 
179
- with gr.Blocks(css=css) as demo:
180
  shuffled_items = gr.State()
181
  title = gr.HTML(
182
  '''<h1>LoRA Roulette 🎰</h1>
@@ -245,5 +258,11 @@ with gr.Blocks(css=css) as demo:
245
  thumbs_down.click(save_preferences, inputs=[lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, output_image, gr.State("down"), last_used_seed], outputs=[thumbs_down, thumbs_down_clicked, thumbs_up])
246
  share_button.click(None, [], [], _js=share_js)
247
 
248
- demo.queue(concurrency_count=2)
249
- demo.launch()
 
 
 
 
 
 
 
14
  import gc
15
  import pickle
16
  import spaces
17
+ import user_history
18
 
19
  lora_list = hf_hub_download(repo_id="multimodalart/LoraTheExplorer", filename="sdxl_loras.json", repo_type="space")
20
 
 
94
  original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
95
 
96
  @spaces.GPU
97
+ def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lora_2_scale=0.5, seed=-1, profile: gr.OAuthProfile | None=None):
98
 
99
  repo_id_1 = shuffled_items[0]['repo']
100
  repo_id_2 = shuffled_items[1]['repo']
 
135
  seed = random.randint(0, 2147483647)
136
  generator = torch.Generator(device="cuda").manual_seed(seed)
137
  image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=20, width=768, height=768, generator=generator).images[0]
138
+
139
+ # save generated images (if logged in)
140
+ user_history.save_image(label=prompt, image=image, profile=profile, metadata={
141
+ "prompt": prompt,
142
+ "negative_prompt": negative_prompt,
143
+ "lora_1_repo_id": shuffled_items[0]['repo'],
144
+ "lora_2_repo_id": shuffled_items[1]['repo'],
145
+ "lora_1_scale": lora_1_scale,
146
+ "lora_2_scale": lora_2_scale,
147
+ "seed": seed,
148
+ })
149
+
150
  return image, gr.update(visible=True), seed, gr.update(visible=True, interactive=True), gr.update(visible=False), gr.update(visible=True, interactive=True), gr.update(visible=False)
151
 
152
  def get_description(item):
 
189
  def hide_post_gen_info():
190
  return gr.update(visible=False)
191
 
192
+ with gr.Blocks() as demo:
193
  shuffled_items = gr.State()
194
  title = gr.HTML(
195
  '''<h1>LoRA Roulette 🎰</h1>
 
258
  thumbs_down.click(save_preferences, inputs=[lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, output_image, gr.State("down"), last_used_seed], outputs=[thumbs_down, thumbs_down_clicked, thumbs_up])
259
  share_button.click(None, [], [], _js=share_js)
260
 
261
+ with gr.Blocks(css=css) as demo_with_history:
262
+ with gr.Tab("Lora Roulette"):
263
+ demo.render()
264
+ with gr.Tab("Past generations"):
265
+ user_history.render()
266
+
267
+ demo_with_history.queue(concurrency_count=2)
268
+ demo_with_history.launch()
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]