Update storage

utils/constants.py

@@ -67,11 +67,6 @@ SCALE_FACTOR = (12/5)
 TMPDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMPDIR, exist_ok=True)
 
-model_extensions = {".glb", ".gltf", ".obj", ".ply"}
-model_extensions_list = list(model_extensions)
-image_extensions = {".png", ".jpg", ".jpeg", ".webp"}
-image_extensions_list = list(image_extensions)
-upload_file_types = model_extensions_list + image_extensions_list
 
 
 PROMPTS = {
@@ -678,5 +673,17 @@ card_colors_alternating = [
 ]
 
 # Constants for URL shortener
-HF_REPO_ID = "Surn/Storage"  # Or your desired repository
-SHORTENER_JSON_FILE = "shortener.json"
+HF_REPO_ID = os.getenv("HF_REPO_ID")
+if not HF_REPO_ID:
+    HF_REPO_ID = "Surn/Storage"  # Replace with your Hugging Face repository ID
+SHORTENER_JSON_FILE = "shortener.json"
+
+model_extensions = {".glb", ".gltf", ".obj", ".ply"}
+model_extensions_list = list(model_extensions)
+image_extensions = {".png", ".jpg", ".jpeg", ".webp"}
+image_extensions_list = list(image_extensions)
+audio_extensions = {".mp3", ".wav", ".ogg", ".flac", ".aac"}
+audio_extensions_list = list(audio_extensions)
+video_extensions = {".mp4"}
+video_extensions_list = list(video_extensions)
+upload_file_types = model_extensions_list + image_extensions_list + audio_extensions_list + video_extensions_list

utils/storage.md

@@ -4,6 +4,8 @@ The `storage.py` module provides helper functions for:
 - Generating permalinks for 3D viewer projects.
 - Uploading files in batches to a Hugging Face repository.
 - Managing URL shortening by storing (short URL, full URL) pairs in a JSON file on the repository.
+- Retrieving full URLs from short URL IDs and vice versa.
+- Handle specific file types for 3D models, images, video and audio.
 
 ## Key Functions
 
@@ -53,7 +55,7 @@ files_for_permalink = [
     "local/path/to/heightmap.png",
     "local/path/to/image.png"
 ]
-repo_id = "Surn/Storage" # Make sure this is defined, e.g., from constants
+repo_id = "Surn/Storage" # Make sure this is defined, e.g., from constants or environment variables
 folder_name = "my_new_model_with_permalink"
 
 upload_result = upload_files_to_repo(

utils/storage.py

@@ -1,5 +1,5 @@
 # utils/storage.py
-__version__ = "0.1.0" # Added version
+__version__ = "0.1.1" # Added version
 import os
 import urllib.parse
 import tempfile
@@ -8,7 +8,7 @@ import json
 import base64
 from huggingface_hub import login, upload_folder, hf_hub_download, HfApi
 from huggingface_hub.utils import RepositoryNotFoundError, EntryNotFoundError
-from utils.constants import HF_API_TOKEN, upload_file_types, model_extensions, image_extensions, HF_REPO_ID, SHORTENER_JSON_FILE
+from utils.constants import HF_API_TOKEN, upload_file_types, model_extensions, image_extensions, audio_extensions, video_extensions, HF_REPO_ID, SHORTENER_JSON_FILE
 from typing import Any, Dict, List, Tuple, Union
 
 # see storage.md for detailed information about the storage module and its functions.
@@ -21,6 +21,8 @@ def generate_permalink(valid_files, base_url_external, permalink_viewer_url="sur
     """
     model_link = None
     images_links = []
+    audio_links = []
+    video_links = []
     for f in valid_files:
         filename = os.path.basename(f)
         ext = os.path.splitext(filename)[1].lower()
@@ -29,6 +31,10 @@ def generate_permalink(valid_files, base_url_external, permalink_viewer_url="sur
                 model_link = f"{base_url_external}/{filename}"
         elif ext in image_extensions:
             images_links.append(f"{base_url_external}/{filename}")
+        elif ext in audio_extensions:
+            audio_links.append(f"{base_url_external}/{filename}")
+        elif ext in video_extensions:
+            video_links.append(f"{base_url_external}/{filename}")
     if model_link and len(images_links) == 2:
         # Construct a permalink to the viewer project with query parameters.
         permalink_viewer_url = f"https://{permalink_viewer_url}/"