Update src/app_hf_space_optimized.py

src/app_hf_space_optimized.py

@@ -22,7 +22,15 @@ import shutil
 import traceback
 import psutil # For memory stats
 
-# Use a more explicit title indicating vertical format
+# --- Set environment variable for Hugging Face cache ---
+# This tells Hugging Face libraries where to store downloaded models
+# /tmp is typically writable on HF Spaces
+HF_CACHE_DIR = "/tmp/hf_cache"
+os.environ['HUGGINGFACE_HUB_CACHE'] = HF_CACHE_DIR
+# Create the cache directory if it doesn't exist
+os.makedirs(HF_CACHE_DIR, exist_ok=True)
+
+
 st.set_page_config(layout="wide", page_title="POV Vertical Video Gen (HF Space)")
 
 # --- Configuration ---
@@ -48,16 +56,19 @@ def display_memory_usage():
     """Displays CPU and GPU memory usage in the sidebar."""
     try:
         process = psutil.Process(os.getpid())
-        cpu_mem = process.memory_info().rss / (1024 * 1024) # MB
-        gpu_mem_info = "N/A"
+        cpu_mem = process.memory_info().rss / (1600 * 1024) # MB (approx max free tier RAM)
+        cpu_mem_percent = (process.memory_info().rss / (1600 * 1024 * 1024)) * 100 # Percentage of 1.6GB
+        cpu_text = f"🧠 CPU Mem: {cpu_mem:.0f} MB ({cpu_mem_percent:.0f}%)"
+
+        gpu_mem_info = "N/A (No GPU)"
         if torch.cuda.is_available():
             # Get current and peak allocated memory
             allocated = torch.cuda.memory_allocated(0) / (1024 * 1024) # MB
-            # reserved = torch.cuda.memory_reserved(0) / (1024 * 1024) # MB # Reserved is less critical than allocated/peak
             peak_allocated = torch.cuda.max_memory_allocated(0) / (1024 * 1024) # MB
             total = torch.cuda.get_device_properties(0).total_memory / (1024 * 1024) # MB
-            gpu_mem_info = f"Alloc: {allocated:.0f}MB | Peak Alloc: {peak_allocated:.0f}MB | Total: {total:.0f}MB"
-        mem_info_placeholder.info(f"🧠 CPU Mem: {cpu_mem:.0f} MB\n⚡ GPU Mem: {gpu_mem_info}")
+            gpu_mem_info = f"⚡ GPU Mem: Alloc: {allocated:.0f}MB | Peak Alloc: {peak_allocated:.0f}MB | Total: {total:.0f}MB"
+
+        mem_info_placeholder.info(f"{cpu_text}\n{gpu_mem_info}")
     except Exception as e:
         mem_info_placeholder.warning(f"Could not get memory info: {e}")
 
@@ -602,7 +613,7 @@ def run_compose_step_ffmpeg(video_results, audio_results, temp_dir, title="final
                 # Specify the input base directory for relative paths using -i
                 (
                     ffmpeg
-                    .input(f"concat:{concat_video_list_path}", format='concat', safe=0, fflags='+igndts', enable_cuda=True) # Use concat protocol, add cuda if available
+                    .input(f"concat:{concat_video_list_path}", format='concat', safe=0, fflags='+igndts')
                     .output(long_video_path, c='copy') # Use stream copy for speed - assumes inputs are compatible codecs/formats
                     .global_args('-hide_banner', '-loglevel', 'error') # Suppress verbose output, show only errors
                     # Run the command from the temporary directory for relative paths to work
@@ -738,7 +749,7 @@ with st.sidebar:
 
     # Add warning about scene count on CPU
     if device == "cpu" and num_scenes_req > 2: # Arbitrary threshold, but 3+ scenes are very risky on CPU
-         st.warning(f"⚠️ Generating {num_scenes_req} scenes on CPU is highly likely to fail due to resource limits.")
+         st.warning(f"⚠️ Requesting {num_scenes_req} scenes on CPU is highly likely to exceed resource limits and fail.")
 
 
     # Display target dimensions and duration clearly
@@ -760,14 +771,14 @@ with st.sidebar:
         st.session_state.num_scenes = num_scenes_req # Store the user-requested number of scenes
         cleanup_temp_dir() # Clean old files before starting a new run
         get_temp_dir() # Ensure a new temp dir path is set for this run
-        st.experimental_rerun() # Trigger a rerun to enter the generation loop
+        st.rerun() # *** UPDATED *** Trigger a rerun to enter the generation loop
 
     st.header("⚠️ Actions")
     # Reset workflow button - disabled if generation is in progress
     if st.button("🔁 Reset Workflow", disabled=st.session_state.generation_in_progress):
         init_state() # Reset all session state
         cleanup_temp_dir() # Also clean files on reset
-        st.experimental_rerun() # Rerun to update UI state and exit generation loop
+        st.rerun() # *** UPDATED *** Rerun to update UI state and exit generation loop
 
     # Clean temp files button - disabled if generation is in progress
     cleanup_button_help = f"Removes files in: {st.session_state.get('temp_dir_path', 'N/A')}"
@@ -875,7 +886,7 @@ if st.session_state.generation_in_progress:
                  # The progress bar update to 100% for 'done' is handled above based on state
              # Trigger a rerun. Streamlit will reload the script, and the logic will continue
              # from the new st.session_state.current_step.
-             st.experimental_rerun()
+             st.rerun() # *** UPDATED *** Trigger a rerun
 
         # If next_step is the same as current_step, it implies an error occurred *within* the step's
         # execution that set next_step to "error", and the logic above decided not to rerun.
@@ -893,7 +904,7 @@ if st.session_state.generation_in_progress:
          st.session_state.current_step = "error" # Set state to error
          st.session_state.generation_in_progress = False # Stop generation
          progress_bar.progress(0) # Reset progress bar on error
-         st.experimental_rerun() # Rerun to show the error state UI and stop execution flow
+         st.rerun() # *** UPDATED *** Rerun to show the error state UI and stop execution flow
 
 
 # --- Display Final Output ---