Add @spaces.GPU decorator for ZeroGPU support

app.py

@@ -10,8 +10,8 @@ import os
 os.environ['OMP_NUM_THREADS'] = '1'
 os.environ['MKL_NUM_THREADS'] = '1'
 os.environ['NUMEXPR_NUM_THREADS'] = '1'
-# Fix CUDA memory fragmentation (use new variable name)
-os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True'
+# Fix CUDA memory fragmentation
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 
 import warnings
 import gradio as gr
@@ -34,6 +34,14 @@ except ImportError:
 # Check if running on Hugging Face Spaces
 HF_SPACE = os.environ.get("SPACE_ID") is not None
 
+# Import spaces module for ZeroGPU support (required for GPU allocation)
+try:
+    import spaces
+    SPACES_AVAILABLE = True
+except ImportError:
+    SPACES_AVAILABLE = False
+    print("⚠️  spaces module not available. GPU decorator will be skipped.")
+
 # 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"
@@ -134,8 +142,19 @@ else:
     print("⚠️  CUDA not available. Make sure GPU hardware is selected in Space settings.")
     print(f"   PyTorch Version: {torch.__version__}")
 
-def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path):
-    """Generate image from text prompt."""
+# Apply @spaces.GPU decorator if available (required for ZeroGPU)
+if SPACES_AVAILABLE and hasattr(spaces, 'GPU'):
+    @spaces.GPU
+    def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path):
+        """Generate image from text prompt."""
+        return _generate_image_impl(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path)
+else:
+    def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path):
+        """Generate image from text prompt."""
+        return _generate_image_impl(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path)
+
+def _generate_image_impl(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path):
+    """Generate image from text prompt (implementation)."""
     # Get checkpoint path (from upload, local, or Model Hub)
     status_msg = ""
     
@@ -288,8 +307,19 @@ def generate_image(prompt, aspect_ratio, cfg, seed, checkpoint_file, config_path
         return None, f"Error: {str(e)}"
 
 
-def generate_video(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image):
-    """Generate video from text prompt."""
+# Apply @spaces.GPU decorator if available (required for ZeroGPU)
+if SPACES_AVAILABLE and hasattr(spaces, 'GPU'):
+    @spaces.GPU
+    def generate_video(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image):
+        """Generate video from text prompt."""
+        return _generate_video_impl(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image)
+else:
+    def generate_video(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image):
+        """Generate video from text prompt."""
+        return _generate_video_impl(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image)
+
+def _generate_video_impl(prompt, aspect_ratio, cfg, seed, target_length, checkpoint_file, config_path, input_image):
+    """Generate video from text prompt (implementation)."""
     # 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