File size: 4,472 Bytes
7cb0b54 d8c5068 7cb0b54 d8c5068 7cb0b54 d8c5068 7cb0b54 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
# modules/storage.py
import os
import urllib.parse
import tempfile
import shutil
from huggingface_hub import login, upload_folder
from modules.constants import HF_API_TOKEN, upload_file_types, model_extensions, image_extensions
def upload_files_to_repo(files, repo_id, folder_name, create_permalink=False, repo_type="dataset"):
"""
Uploads multiple files to a Hugging Face repository using a batch upload approach via upload_folder.
Parameters:
files (list): A list of file paths (str) to upload.
repo_id (str): The repository ID on Hugging Face for storage, e.g. "Surn/Storage".
folder_name (str): The subfolder within the repository where files will be saved.
create_permalink (bool): If True and if exactly three files are uploaded (1 model and 2 images),
returns a single permalink to the project with query parameters.
Otherwise, returns individual permalinks for each file.
repo_type (str): Repository type ("space", "dataset", etc.). Default is "dataset".
Returns:
If create_permalink is True and files match the criteria:
tuple: (response, permalink) where response is the output of the batch upload
and permalink is the URL string (with fully qualified file paths) for the project.
Otherwise:
list: A list of tuples (response, permalink) for each file.
"""
# Log in using the HF API token.
login(token=HF_API_TOKEN)
valid_files = []
# Ensure folder_name does not have a trailing slash.
folder_name = folder_name.rstrip("/")
# Filter for valid files based on allowed extensions.
for f in files:
file_name = f if isinstance(f, str) else f.name if hasattr(f, "name") else None
if file_name is None:
continue
ext = os.path.splitext(file_name)[1].lower()
if ext in upload_file_types:
valid_files.append(f)
if not valid_files:
return [] # or raise an exception
# Create a temporary directory; copy valid files directly into it (without using folder_name).
with tempfile.TemporaryDirectory() as temp_dir:
for file_path in valid_files:
filename = os.path.basename(file_path)
dest_path = os.path.join(temp_dir, filename)
shutil.copy(file_path, dest_path)
# Batch upload all files in the temporary folder.
# Files will be uploaded under the folder (path_in_repo) given by folder_name.
response = upload_folder(
folder_path=temp_dir,
repo_id=repo_id,
repo_type=repo_type,
path_in_repo=folder_name,
commit_message="Batch upload files"
)
# Construct external URLs for each uploaded file.
# For datasets, files are served at:
# https://huggingface.co/datasets/<repo_id>/resolve/main/<folder_name>/<filename>
base_url_external = f"https://huggingface.co/datasets/{repo_id}/resolve/main/{folder_name}"
individual_links = []
for file_path in valid_files:
filename = os.path.basename(file_path)
link = f"{base_url_external}/{filename}"
individual_links.append(link)
# If exactly 3 files are provided and create_permalink is True,
# check if they contain 1 model file and 2 image files.
if create_permalink and len(valid_files) == 3:
model_link = None
images_links = []
for f in valid_files:
filename = os.path.basename(f)
ext = os.path.splitext(filename)[1].lower()
if ext in model_extensions:
if model_link is None:
model_link = f"{base_url_external}/{filename}"
elif ext in image_extensions:
images_links.append(f"{base_url_external}/{filename}")
if model_link and len(images_links) == 2:
# Construct a permalink to the viewer project with querystring parameters.
base_viewer_url = "https://huggingface.co/spaces/Surn/3D-Viewer"
params = {"3d": model_link, "hm": images_links[0], "image": images_links[1]}
query_str = urllib.parse.urlencode(params)
permalink = f"{base_viewer_url}?{query_str}"
return response, permalink
# Otherwise, return individual tuples for each file.
return [(response, link) for link in individual_links]
|