gradio/scripts/upload_demo_to_space.py
pngwn f61d7024d9
chunk space uploads (#9360)
* chunk space uploads

* Update upload_demo_to_space.py

Co-authored-by: Lucain <lucainp@gmail.com>

* address comments + tweak CI

---------

Co-authored-by: Lucain <lucainp@gmail.com>
2024-09-17 15:21:49 +08:00

149 lines
4.8 KiB
Python

import argparse
import pathlib
import shutil
import sys
import tempfile
import textwrap
from typing import Optional
from huggingface_hub import CommitOperationAdd, HfApi
def upload_demo_to_space(
demo_name: str, space_id: str, hf_token: str, gradio_version: Optional[str]
):
"""
Upload a demo from the demo directory to a Hugging Face Space in chunks of 50 files per commit.
Args:
demo_name: The name of the demo to upload.
space_id: The ID of the space to upload the demo to (e.g., username/space_name).
hf_token: Hugging Face API token with write permissions to the space.
gradio_version: If provided, sets the Gradio version in the created space.
Returns:
str: URL of the uploaded Hugging Face Space.
"""
print(f"Uploading demo '{demo_name}' to space '{space_id}'...")
def split_into_chunks(lst: list, n: int) -> list[list]:
for i in range(0, len(lst), n):
yield lst[i : i + n]
api = HfApi()
print("Creating repository...")
# Create the repository if it doesn't exist
space_url = api.create_repo(
repo_id=space_id,
space_sdk="gradio",
repo_type="space",
token=hf_token,
exist_ok=True,
)
space_id = space_url.repo_id
with tempfile.TemporaryDirectory() as tmpdir:
demo_path = pathlib.Path.cwd() / "demo" / demo_name
if not demo_path.exists():
raise FileNotFoundError(f"Demo path '{demo_path}' does not exist.")
shutil.copytree(demo_path, tmpdir, dirs_exist_ok=True)
# update README.md with Gradio version if provided
if gradio_version:
readme = pathlib.Path(tmpdir, "README.md")
readme_content = f"""
---
title: {space_id.split("/")[-1]}
emoji: 💩
colorFrom: indigo
colorTo: indigo
sdk: gradio
sdk_version: {gradio_version}
app_file: run.py
pinned: false
---
"""
readme.write_text(textwrap.dedent(readme_content))
print("Uploading files to Hugging Face Space...")
# Create the repository if it doesn't exist
api.create_repo(
repo_id=space_id,
space_sdk="gradio",
repo_type="space",
token=hf_token,
exist_ok=True,
)
print("Uploading files...")
all_files = sorted([p for p in pathlib.Path(tmpdir).rglob("*") if p.is_file()])
relative_files = [p.relative_to(tmpdir) for p in all_files]
# Create CommitOperationAdd objects for all files
operations = [
CommitOperationAdd(
path_in_repo=str(rel_path).replace("\\", "/"),
path_or_fileobj=str(pathlib.Path(tmpdir) / rel_path),
)
for rel_path in relative_files
]
# Split operations into chunks of 50
operation_chunks = list(split_into_chunks(operations, 50))
for idx, chunk in enumerate(operation_chunks, start=1):
commit_message = f"Commit {idx}: Add {len(chunk)} file(s)"
try:
api.create_commit(
repo_id=space_id,
operations=chunk,
commit_message=commit_message,
token=hf_token,
repo_type="space",
)
print(f"Successfully committed chunk {idx} with {len(chunk)} file(s).")
except Exception as e:
print(f"Failed to commit chunk {idx}: {e}")
raise e
return f"https://huggingface.co/spaces/{space_id}"
if __name__ == "__main__":
print("Starting upload...")
parser = argparse.ArgumentParser(
description="Upload a demo to a Hugging Face Space in chunks."
)
parser.add_argument("demo_name", type=str, help="Name of the demo to upload")
parser.add_argument(
"space_id",
type=str,
help="ID of the space to upload the demo to (e.g., username/space_name)",
)
parser.add_argument("hf_token", type=str, help="Hugging Face API token")
parser.add_argument(
"--gradio-version",
type=str,
help="If provided, sets the Gradio version in the created space to the given version.",
)
args = parser.parse_args()
try:
print("Uploading demo to Hugging Face Space...")
new_space_url = upload_demo_to_space(
args.demo_name, args.space_id, args.hf_token, args.gradio_version
)
print("Demo successfully uploaded to:")
# do not change the final print statement,
# it must appear as the last line of the logs for CI to pick up the URL
print(new_space_url)
except Exception as error:
print(f"An error occurred during upload: {error}")
sys.exit(1)