Fix checkpoint download logic to properly handle Model Hub downloads

app.py

@@ -45,10 +45,18 @@ VIDEO_CHECKPOINT_REPO = "GlobalStudio/starflow-v-7b-checkpoint"  # Update this a
 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 checkpoint_file is not None and checkpoint_file != "":
         if hasattr(checkpoint_file, 'name'):
             return checkpoint_file.name
-        return str(checkpoint_file)
+        checkpoint_str = str(checkpoint_file)
+        # If it's a file path that exists, use it
+        if os.path.exists(checkpoint_str):
+            return checkpoint_str
+        # If it's the default path but doesn't exist, continue to download
+        if checkpoint_str == default_local_path and not os.path.exists(checkpoint_str):
+            pass  # Continue to download logic below
+        elif checkpoint_str != default_local_path:
+            return checkpoint_str
     
     # Try local path first
     if os.path.exists(default_local_path):
@@ -108,6 +116,13 @@ else:
 def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path):
     """Generate image from text prompt."""
     # Get checkpoint path (from upload, local, or Model Hub)
+    status_msg = ""
+    
+    # Handle checkpoint file (might be string from hidden Textbox)
+    if checkpoint_file == DEFAULT_IMAGE_CHECKPOINT or checkpoint_file == "" or checkpoint_file is None:
+        # Use Model Hub download
+        checkpoint_file = None
+    
     result = get_checkpoint_path(
         checkpoint_file, 
         DEFAULT_IMAGE_CHECKPOINT,
@@ -121,10 +136,10 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
     checkpoint_path = result
     
     # Show status
-    status_msg = f"Using checkpoint: {checkpoint_path}\n"
+    status_msg += f"Using checkpoint: {checkpoint_path}\n"
     
     if not os.path.exists(checkpoint_path):
-        return None, f"Error: Checkpoint file not found at {checkpoint_path}.\n\nPlease verify:\n1. Model Hub repo exists: {IMAGE_CHECKPOINT_REPO}\n2. File name matches: starflow_3B_t2i_256x256.pth\n3. Repo is Public (not Private)"
+        return None, f"Error: Checkpoint file not found at {checkpoint_path}.\n\nPlease verify:\n1. Model Hub repo exists: {IMAGE_CHECKPOINT_REPO}\n2. File name matches: starflow_3B_t2i_256x256.pth\n3. Repo is Public (not Private)\n\nAttempting to download from Model Hub..."
     
     if not config_path or not os.path.exists(config_path):
         return None, "Error: Config file not found. Please ensure config file exists."
@@ -219,6 +234,11 @@ 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."""
+    # Handle checkpoint file (might be string from hidden Textbox)
+    if checkpoint_file == DEFAULT_VIDEO_CHECKPOINT or checkpoint_file == "" or checkpoint_file is None:
+        # Use Model Hub download
+        checkpoint_file = None
+    
     # Get checkpoint path (from upload, local, or Model Hub)
     result = get_checkpoint_path(
         checkpoint_file,