Add Model Hub support for downloading checkpoints (solves 1GB storage limit)

app.py

@@ -12,6 +12,14 @@ import subprocess
 import pathlib
 from pathlib import Path
 
+# Try to import huggingface_hub for downloading checkpoints
+try:
+    from huggingface_hub import hf_hub_download
+    HF_HUB_AVAILABLE = True
+except ImportError:
+    HF_HUB_AVAILABLE = False
+    print("⚠️  huggingface_hub not available. Install with: pip install huggingface_hub")
+
 # Check if running on Hugging Face Spaces
 HF_SPACE = os.environ.get("SPACE_ID") is not None
 
@@ -19,6 +27,41 @@ HF_SPACE = os.environ.get("SPACE_ID") is not None
 DEFAULT_IMAGE_CHECKPOINT = "ckpts/starflow_3B_t2i_256x256.pth"
 DEFAULT_VIDEO_CHECKPOINT = "ckpts/starflow-v_7B_t2v_caus_480p_v3.pth"
 
+# Model Hub repositories (if using Hugging Face Model Hub)
+# Set these to your Model Hub repo IDs if you upload checkpoints there
+IMAGE_CHECKPOINT_REPO = None  # e.g., "GlobalStudio/starflow-3b-checkpoint"
+VIDEO_CHECKPOINT_REPO = None   # e.g., "GlobalStudio/starflow-v-7b-checkpoint"
+
+def get_checkpoint_path(checkpoint_file, default_local_path, repo_id=None, filename=None):
+    """Get checkpoint path, downloading from Hub if needed."""
+    # If user uploaded a file, use it
+    if checkpoint_file is not None:
+        if hasattr(checkpoint_file, 'name'):
+            return checkpoint_file.name
+        return str(checkpoint_file)
+    
+    # Try local path first
+    if os.path.exists(default_local_path):
+        return default_local_path
+    
+    # Try downloading from Model Hub if configured
+    if repo_id and filename and HF_HUB_AVAILABLE:
+        try:
+            print(f"📥 Downloading checkpoint from {repo_id}...")
+            checkpoint_path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                cache_dir="/tmp/checkpoints",
+                local_files_only=False
+            )
+            print(f"✅ Checkpoint downloaded to: {checkpoint_path}")
+            return checkpoint_path
+        except Exception as e:
+            return None, f"Error downloading checkpoint: {str(e)}"
+    
+    # No checkpoint found
+    return None, f"Checkpoint not found. Please upload a checkpoint file or configure Model Hub repository."
+
 # Verify CUDA availability (will be True on HF Spaces with GPU hardware)
 if torch.cuda.is_available():
     print(f"✅ CUDA available! Device: {torch.cuda.get_device_name(0)}")
@@ -29,19 +72,18 @@ else:
 
 def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path):
     """Generate image from text prompt."""
-    # Use uploaded file if provided, otherwise try default persistent checkpoint
-    if checkpoint_file is None:
-        # Try to use pre-uploaded checkpoint
-        checkpoint_path = DEFAULT_IMAGE_CHECKPOINT
-        if not os.path.exists(checkpoint_path):
-            return None, f"Error: No checkpoint found. Please upload a checkpoint file or ensure '{DEFAULT_IMAGE_CHECKPOINT}' exists in Space Files."
-    else:
-        # Handle Gradio file object (user uploaded)
-        if hasattr(checkpoint_file, 'name'):
-            checkpoint_path = checkpoint_file.name
-        else:
-            checkpoint_path = str(checkpoint_file)
+    # Get checkpoint path (from upload, local, or Model Hub)
+    result = get_checkpoint_path(
+        checkpoint_file, 
+        DEFAULT_IMAGE_CHECKPOINT,
+        IMAGE_CHECKPOINT_REPO,
+        "starflow_3B_t2i_256x256.pth"
+    )
+    
+    if isinstance(result, tuple) and result[0] is None:
+        return None, result[1]  # Error message
     
+    checkpoint_path = result
     if not os.path.exists(checkpoint_path):
         return None, f"Error: Checkpoint file not found at {checkpoint_path}."
     
@@ -91,19 +133,18 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
 
 def generate_video(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image):
     """Generate video from text prompt."""
-    # Use uploaded file if provided, otherwise try default persistent checkpoint
-    if checkpoint_file is None:
-        # Try to use pre-uploaded checkpoint
-        checkpoint_path = DEFAULT_VIDEO_CHECKPOINT
-        if not os.path.exists(checkpoint_path):
-            return None, f"Error: No checkpoint found. Please upload a checkpoint file or ensure '{DEFAULT_VIDEO_CHECKPOINT}' exists in Space Files."
-    else:
-        # Handle Gradio file object (user uploaded)
-        if hasattr(checkpoint_file, 'name'):
-            checkpoint_path = checkpoint_file.name
-        else:
-            checkpoint_path = str(checkpoint_file)
+    # Get checkpoint path (from upload, local, or Model Hub)
+    result = get_checkpoint_path(
+        checkpoint_file,
+        DEFAULT_VIDEO_CHECKPOINT,
+        VIDEO_CHECKPOINT_REPO,
+        "starflow-v_7B_t2v_caus_480p_v3.pth"
+    )
+    
+    if isinstance(result, tuple) and result[0] is None:
+        return None, result[1]  # Error message
     
+    checkpoint_path = result
     if not os.path.exists(checkpoint_path):
         return None, f"Error: Checkpoint file not found at {checkpoint_path}."
     

requirements_hf.txt

@@ -9,6 +9,7 @@ torchvision>=0.15.0
 # Core dependencies
 transformers>=4.30.0
 accelerate>=0.20.0
+huggingface_hub>=0.20.0
 torchinfo
 einops
 scipy