gradio/demo/state_cleanup/run.py
Freddy Boulton 6a4bf7abe2
Delete user state when they close the tab. Add an unload event for the demo and a delete_callback on gr.State to let developers control how resources are cleaned up (#7829)
* Delete state

* add changeset

* Delete state

* WIP

* Add load event

* Working ttl

* unload e2e test

* Clean up

* add changeset

* Fix notebook

* add changeset

* Connect to heartbeat in python client

* 15 second heartbeat

* Demo for unload

* Add notebook

* add changeset

* Fix docs

* revert demo changes

* Add docstrings

* lint 🙄

* Edit

* handle shutdown issue

* state comments

* client test

* Fix:

* Fix e2e test

* 3.11 incompatibility

* delete after one hour

* lint + highlight

* Update .changeset/better-tires-shave.md

Co-authored-by: Abubakar Abid <abubakar@huggingface.co>

* Update .changeset/better-tires-shave.md

Co-authored-by: Abubakar Abid <abubakar@huggingface.co>

---------

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>
Co-authored-by: Abubakar Abid <abubakar@huggingface.co>
2024-04-01 18:31:56 -04:00

53 lines
1.8 KiB
Python

from __future__ import annotations
import gradio as gr
import numpy as np
from PIL import Image
from pathlib import Path
import secrets
import shutil
current_dir = Path(__file__).parent
def generate_random_img(history: list[Image.Image], request: gr.Request):
"""Generate a random red, green, blue, orange, yellor or purple image."""
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 165, 0), (255, 255, 0), (128, 0, 128)]
color = colors[np.random.randint(0, len(colors))]
img = Image.new('RGB', (100, 100), color)
user_dir: Path = current_dir / request.username # type: ignore
user_dir.mkdir(exist_ok=True)
path = user_dir / f"{secrets.token_urlsafe(8)}.webp"
img.save(path)
history.append(img)
return img, history, history
def delete_directory(req: gr.Request):
if not req.username:
return
user_dir: Path = current_dir / req.username
shutil.rmtree(str(user_dir))
with gr.Blocks() as demo:
gr.Markdown("""# State Cleanup Demo
🖼️ Images are saved in a user-specific directory and deleted when the users closes the page via demo.unload.
""")
with gr.Row():
with gr.Column(scale=1):
with gr.Row():
img = gr.Image(label="Generated Image", height=300, width=300)
with gr.Row():
gen = gr.Button(value="Generate")
with gr.Row():
history = gr.Gallery(label="Previous Generations", height=500, columns=10)
state = gr.State(value=[], delete_callback=lambda v: print("STATE DELETED"))
demo.load(generate_random_img, [state], [img, state, history])
gen.click(generate_random_img, [state], [img, state, history])
demo.unload(delete_directory)
demo.launch(auth=lambda user,pwd: True,
auth_message="Enter any username and password to continue")