Update storage 0.1.1

modules/constants.py

@@ -27,7 +27,11 @@ 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
+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
 
 default_slider_images = [
             "images/slider/beeuty_545jlbh1_v12_alpha96_300dpi.png",
@@ -37,7 +41,9 @@ default_slider_images = [
 default_model_3d = "models/beeuty_545jlbh1_300dpi.glb"
 
 # Constants for URL shortener
-HF_REPO_ID = "Surn/Storage"  # Replace with your Hugging Face repository ID
+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"  # The name of your JSON file in the repo
 
 # Open Graph Defaults

modules/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(

modules/storage.py

@@ -1,5 +1,5 @@
 # modules/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 modules.constants import HF_API_TOKEN, upload_file_types, model_extensions, image_extensions, HF_REPO_ID, SHORTENER_JSON_FILE
+from modules.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}/"
@@ -212,7 +218,7 @@ def _upload_json_to_repo(data, repo_id, json_file_name, repo_type="dataset"):
         temp_dir_for_json = os.getenv("TMPDIR", tempfile.gettempdir())
         os.makedirs(temp_dir_for_json, exist_ok=True)
 
-        with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json", dir=temp_dir_for_json) as tmp_file:        
+        with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json", dir=temp_dir_for_json) as tmp_file:
             json.dump(data, tmp_file, indent=2)
             tmp_file_path = tmp_file.name
         
@@ -251,7 +257,7 @@ def _find_url_in_json(data, short_url=None, full_url=None):
 
 def _add_url_to_json(data, short_url, full_url):
     """Adds a new short_url/full_url pair to the data. Returns updated data."""
-    if data is None:
+    if data is None: 
         data = []
     data.append({"short_url": short_url, "full_url": full_url})
     return data