Spaces:

Wauplin
/

space_to_dataset_saver

Running

App Files Files Community

space_to_dataset_saver / app_parquet.py

Wauplin HF staff

Update app_parquet.py

70b2563 9 months ago

raw history blame contribute delete

No virus

9.09 kB

	# Taken from https://huggingface.co/spaces/hysts-samples/save-user-preferences
	# Credits to @@hysts
	import datetime
	import json
	import shutil
	import tempfile
	import uuid
	from pathlib import Path
	from typing import Any, Dict, List, Optional, Union

	import gradio as gr
	import pyarrow as pa
	import pyarrow.parquet as pq
	from gradio_client import Client
	from huggingface_hub import CommitScheduler
	from huggingface_hub.hf_api import HfApi

	#######################
	# Parquet scheduler #
	# Run in scheduler.py #
	#######################


	class ParquetScheduler(CommitScheduler):
	"""
	Usage: configure the scheduler with a repo id. Once started, you can add data to be uploaded to the Hub. 1 `.append`
	call will result in 1 row in your final dataset.

	```py
	# Start scheduler
	>>> scheduler = ParquetScheduler(repo_id="my-parquet-dataset")

	# Append some data to be uploaded
	>>> scheduler.append({...})
	>>> scheduler.append({...})
	>>> scheduler.append({...})
	```

	The scheduler will automatically infer the schema from the data it pushes.
	Optionally, you can manually set the schema yourself:

	```py
	>>> scheduler = ParquetScheduler(
	... repo_id="my-parquet-dataset",
	... schema={
	... "prompt": {"_type": "Value", "dtype": "string"},
	... "negative_prompt": {"_type": "Value", "dtype": "string"},
	... "guidance_scale": {"_type": "Value", "dtype": "int64"},
	... "image": {"_type": "Image"},
	... },
	... )

	See https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Value for the list of
	possible values.
	"""

	def __init__(
	self,
	*,
	repo_id: str,
	schema: Optional[Dict[str, Dict[str, str]]] = None,
	every: Union[int, float] = 5,
	path_in_repo: Optional[str] = "data",
	repo_type: Optional[str] = "dataset",
	revision: Optional[str] = None,
	private: bool = False,
	token: Optional[str] = None,
	allow_patterns: Union[List[str], str, None] = None,
	ignore_patterns: Union[List[str], str, None] = None,
	hf_api: Optional[HfApi] = None,
	) -> None:
	super().__init__(
	repo_id=repo_id,
	folder_path="dummy", # not used by the scheduler
	every=every,
	path_in_repo=path_in_repo,
	repo_type=repo_type,
	revision=revision,
	private=private,
	token=token,
	allow_patterns=allow_patterns,
	ignore_patterns=ignore_patterns,
	hf_api=hf_api,
	)

	self._rows: List[Dict[str, Any]] = []
	self._schema = schema

	def append(self, row: Dict[str, Any]) -> None:
	"""Add a new item to be uploaded."""
	with self.lock:
	self._rows.append(row)

	def push_to_hub(self):
	# Check for new rows to push
	with self.lock:
	rows = self._rows
	self._rows = []
	if not rows:
	return
	print(f"Got {len(rows)} item(s) to commit.")

	# Load images + create 'features' config for datasets library
	schema: Dict[str, Dict] = self._schema or {}
	path_to_cleanup: List[Path] = []
	for row in rows:
	for key, value in row.items():
	# Infer schema (for `datasets` library)
	if key not in schema:
	schema[key] = _infer_schema(key, value)

	# Load binary files if necessary
	if schema[key]["_type"] in ("Image", "Audio"):
	# It's an image or audio: we load the bytes and remember to cleanup the file
	file_path = Path(value)
	if file_path.is_file():
	row[key] = {
	"path": file_path.name,
	"bytes": file_path.read_bytes(),
	}
	path_to_cleanup.append(file_path)

	# Complete rows if needed
	for row in rows:
	for feature in schema:
	if feature not in row:
	row[feature] = None

	# Export items to Arrow format
	table = pa.Table.from_pylist(rows)

	# Add metadata (used by datasets library)
	table = table.replace_schema_metadata(
	{"huggingface": json.dumps({"info": {"features": schema}})}
	)

	# Write to parquet file
	archive_file = tempfile.NamedTemporaryFile()
	pq.write_table(table, archive_file.name)

	# Upload
	self.api.upload_file(
	repo_id=self.repo_id,
	repo_type=self.repo_type,
	revision=self.revision,
	path_in_repo=f"{uuid.uuid4()}.parquet",
	path_or_fileobj=archive_file.name,
	)
	print(f"Commit completed.")

	# Cleanup
	archive_file.close()
	for path in path_to_cleanup:
	path.unlink(missing_ok=True)


	def _infer_schema(key: str, value: Any) -> Dict[str, str]:
	"""
	Infer schema for the `datasets` library.

	See https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Value.
	"""
	if "image" in key:
	return {"_type": "Image"}
	if "audio" in key:
	return {"_type": "Audio"}
	if isinstance(value, int):
	return {"_type": "Value", "dtype": "int64"}
	if isinstance(value, float):
	return {"_type": "Value", "dtype": "float64"}
	if isinstance(value, bool):
	return {"_type": "Value", "dtype": "bool"}
	if isinstance(value, bytes):
	return {"_type": "Value", "dtype": "binary"}
	# Otherwise in last resort => convert it to a string
	return {"_type": "Value", "dtype": "string"}


	#################
	# Gradio app #
	# Run in app.py #
	#################

	PARQUET_DATASET_DIR = Path("parquet_dataset")
	PARQUET_DATASET_DIR.mkdir(parents=True, exist_ok=True)

	scheduler = ParquetScheduler(repo_id="example-space-to-dataset-parquet")

	# client = Client("stabilityai/stable-diffusion") # Space is paused
	client = Client("runwayml/stable-diffusion-v1-5")


	def generate(prompt: str) -> tuple[str, list[str]]:
	"""Generate images on 'submit' button."""
	# Generate from https://huggingface.co/spaces/stabilityai/stable-diffusion
	# out_dir = client.predict(prompt, "", 9, fn_index=1) # Space 'stabilityai/stable-diffusion' is paused
	out_dir = client.predict(prompt, fn_index=1)
	with (Path(out_dir) / "captions.json").open() as f:
	paths = list(json.load(f).keys())

	# Save config used to generate data
	with tempfile.NamedTemporaryFile(
	mode="w", suffix=".json", delete=False
	) as config_file:
	json.dump(
	{"prompt": prompt, "negative_prompt": "", "guidance_scale": 9}, config_file
	)

	return config_file.name, paths


	def get_selected_index(evt: gr.SelectData) -> int:
	"""Select "best" image."""
	return evt.index


	def save_preference(
	config_path: str, gallery: list[dict[str, Any]], selected_index: int
	) -> None:
	"""Save preference, i.e. move images to a new folder and send paths+config to scheduler."""
	save_dir = PARQUET_DATASET_DIR / f"{uuid.uuid4()}"
	save_dir.mkdir(parents=True, exist_ok=True)

	# Load config
	with open(config_path) as f:
	data = json.load(f)

	# Add selected item + timestamp
	data["selected_index"] = selected_index
	data["timestamp"] = datetime.datetime.utcnow().isoformat()

	# Copy and add images
	for index, path in enumerate(x["name"] for x in gallery):
	name = f"{index:03d}"
	dst_path = save_dir / f"{name}{Path(path).suffix}"
	shutil.move(path, dst_path)
	data[f"image_{name}"] = dst_path

	# Send to scheduler
	scheduler.append(data)


	def clear() -> tuple[dict, dict, dict]:
	"""Clear all values once saved."""
	return (gr.update(value=None), gr.update(value=None), gr.update(interactive=False))


	def get_demo():
	with gr.Group():
	prompt = gr.Text(show_label=False, placeholder="Prompt")
	config_path = gr.Text(visible=False)
	gallery = gr.Gallery(show_label=False).style(
	columns=2, rows=2, height="600px", object_fit="scale-down"
	)
	selected_index = gr.Number(visible=False, precision=0)
	save_preference_button = gr.Button("Save preference", interactive=False)

	# Generate images on submit
	prompt.submit(fn=generate, inputs=prompt, outputs=[config_path, gallery],).success(
	fn=lambda: gr.update(interactive=True),
	outputs=save_preference_button,
	queue=False,
	)

	# Save preference on click
	gallery.select(
	fn=get_selected_index,
	outputs=selected_index,
	queue=False,
	)
	save_preference_button.click(
	fn=save_preference,
	inputs=[config_path, gallery, selected_index],
	queue=False,
	).then(
	fn=clear,
	outputs=[config_path, gallery, save_preference_button],
	queue=False,
	)