Update app.py to support persistent checkpoints

app.py

@@ -15,6 +15,10 @@ from pathlib import Path
 # Check if running on Hugging Face Spaces
 HF_SPACE = os.environ.get("SPACE_ID") is not None
 
+# Default checkpoint paths (if uploaded to Space Files)
+DEFAULT_IMAGE_CHECKPOINT = "ckpts/starflow_3B_t2i_256x256.pth"
+DEFAULT_VIDEO_CHECKPOINT = "ckpts/starflow-v_7B_t2v_caus_480p_v3.pth"
+
 # 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)}")
@@ -25,14 +29,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:
-        return None, "Error: Please upload a checkpoint file."
-    
-    # Handle Gradio file object
-    if hasattr(checkpoint_file, 'name'):
-        checkpoint_path = checkpoint_file.name
+        # 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:
-        checkpoint_path = str(checkpoint_file)
+        # Handle Gradio file object (user uploaded)
+        if hasattr(checkpoint_file, 'name'):
+            checkpoint_path = checkpoint_file.name
+        else:
+            checkpoint_path = str(checkpoint_file)
     
     if not os.path.exists(checkpoint_path):
         return None, f"Error: Checkpoint file not found at {checkpoint_path}."
@@ -83,14 +91,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:
-        return None, "Error: Please upload a checkpoint file."
-    
-    # Handle Gradio file object
-    if hasattr(checkpoint_file, 'name'):
-        checkpoint_path = checkpoint_file.name
+        # 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:
-        checkpoint_path = str(checkpoint_file)
+        # Handle Gradio file object (user uploaded)
+        if hasattr(checkpoint_file, 'name'):
+            checkpoint_path = checkpoint_file.name
+        else:
+            checkpoint_path = str(checkpoint_file)
     
     if not os.path.exists(checkpoint_path):
         return None, f"Error: Checkpoint file not found at {checkpoint_path}."
@@ -173,7 +185,7 @@ with gr.Blocks(title="STARFlow - Text-to-Image & Video Generation") as demo:
                         lines=3
                     )
                     image_checkpoint = gr.File(
-                        label="Model Checkpoint (.pth file)",
+                        label="Model Checkpoint (.pth file) - Optional if already uploaded to Space",
                         file_types=[".pth"]
                     )
                     image_config = gr.Textbox(
@@ -210,7 +222,7 @@ with gr.Blocks(title="STARFlow - Text-to-Image & Video Generation") as demo:
                         lines=3
                     )
                     video_checkpoint = gr.File(
-                        label="Model Checkpoint (.pth file)",
+                        label="Model Checkpoint (.pth file) - Optional if already uploaded to Space",
                         file_types=[".pth"]
                     )
                     video_config = gr.Textbox(