Lazy load models for ZeroGPU Space startup

app.py

@@ -20,6 +20,17 @@ import html
 import threading
 import traceback
 
+try:
+    import spaces
+except ImportError:
+    class _SpacesNoop:
+        def GPU(self, *args, **kwargs):
+            def decorator(fn):
+                return fn
+            return decorator
+
+    spaces = _SpacesNoop()
+
 # Global variables for controlling the flow of the program
 MAX_LINES_DISPLAYED = 10
 MAX_NUM_OF_SPEAKERS = 8
@@ -30,6 +41,7 @@ FILE_PROCESSING_STOPPED = threading.Event()
 FILE_PROCESSING_STOPPED.set()
 MICROPHONE_PROCESSING_ENABLED: bool = False
 MICROPHONE_PROCESSING_LOCK = threading.Lock()
+GPU_DURATION_SECONDS = int(os.environ.get("SPACES_GPU_DURATION", "120"))
 CONFIG_RENEWED_FLAG: bool = False
 START_TIME_OFFSET_EXAMPLES: float = -1.6 # seconds
 START_TIME_OFFSET_FILE_UPLOAD: float = -0.8 # seconds
@@ -147,13 +159,26 @@ config_file = MultitalkerTranscriptionConfig(
     audio_pre_loading=True,
 )
 
-# Initialize the streaming session
-session_mic= StreamingSession()
-session_mic.setup_streaming_session(config_mic)
+# Initialize streaming session containers. Models are loaded lazily so ZeroGPU
+# Spaces can start the app before CUDA is allocated to a request.
+session_mic = StreamingSession()
 session_file = StreamingSession()
-session_file.setup_streaming_session(config_file)
+MODEL_INIT_LOCK = threading.Lock()
+MODELS_INITIALIZED = False
 MAX_NUM_OF_SPEAKERS = min(config_mic.max_num_of_spks, config_file.max_num_of_spks)
 
+
+def ensure_sessions_initialized(reload=False):
+    """Load or reload ASR/diarization sessions on demand."""
+    global MODELS_INITIALIZED
+    with MODEL_INIT_LOCK:
+        should_reload = reload and MODELS_INITIALIZED
+        if MODELS_INITIALIZED and not reload:
+            return
+        session_mic.setup_streaming_session(config_mic, reload=should_reload)
+        session_file.setup_streaming_session(config_file, reload=should_reload)
+        MODELS_INITIALIZED = True
+
 def render_diarization_state(diar_state, config):
     """Render Sortformer's speaker cache, FIFO, and current chunk predictions."""
     num_speakers = int(config.max_num_of_spks)
@@ -200,10 +225,12 @@ def format_transcript(text):
     if not text:
         return ""
 
-    single_speaker_mode = (
-        session_mic.multispk_asr_streamer._single_speaker_mode
-        or session_file.multispk_asr_streamer._single_speaker_mode
-    )
+    single_speaker_mode = False
+    if MODELS_INITIALIZED:
+        single_speaker_mode = (
+            session_mic.multispk_asr_streamer._single_speaker_mode
+            or session_file.multispk_asr_streamer._single_speaker_mode
+        )
     current_color = "#76B900" if single_speaker_mode else "#A8A8A8"
 
     for raw_line in (line for line in text.strip().splitlines() if line.strip()):
@@ -356,6 +383,7 @@ def update_config_for_file(evt: gr.SelectData, state=None):
 
 def file_processing(filepath, state, file_upload=False, start_time_offset=0.0):
     """Handles file uploads and processes the audio in chunks with real-time simulation."""
+    ensure_sessions_initialized()
     print(f"[DEBUG] file_processing() FIRST LINE: filepath: {filepath}, state[is_playing]: {state.get('is_playing', False)}, file_upload: {file_upload}")
     if not filepath or state.get("should_stop", False):
         yield state, state["diar_html"], format_transcript(state.get("full_transcript", ""))
@@ -499,6 +527,7 @@ async def show_progress_bar_with_sleep(filepath, num_steps=500, seconds_per_mb=0
 
 async def file_processing_example_click_async(filepath, state, start_time_offset=0.0):
     """Async wrapper for file_processing with pause/resume support (doesn't update audio component)."""
+    ensure_sessions_initialized()
     state["should_stop"] = False
     state["is_playing"] = False  
     state["pause_start_time"] = None
@@ -614,12 +643,14 @@ def file_processing_example_click(filepath, state, start_time_offset=0.0):
         loop.close()
         FILE_PROCESSING_STOPPED.set()
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def file_processing_example_click_upload(filepath, state):
     """Wrapper for file upload with upload offset."""
     # file_processing_example_click is a generator, so yield from it
     for result in file_processing_example_click(filepath, state, start_time_offset=START_TIME_OFFSET_FILE_UPLOAD):
         yield result
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def file_processing_example_click_examples(filepath, state):
     """Wrapper for examples with examples offset."""
     # file_processing_example_click is a generator, so yield from it
@@ -665,6 +696,7 @@ async def file_processing_upload_event_async(filepath, state):
         
         yield s, html_diar, transcript
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def file_processing_upload_event(filepath, state): 
     """Wrapper that runs the async function synchronously for Gradio compatibility."""
     print(f"[DEBUG] First line of file_processing_upload_event(): filepath: {filepath}, state[is_playing]: {state.get('is_playing', False)}")
@@ -730,6 +762,7 @@ def example_selection_stop_file_processing(state, filepath):
 
 def file_upload_stop_file_processing(state):
     """Stops file processing when audio is cleared and prepares for new file."""
+    ensure_sessions_initialized()
     state["mode"] = "idle"
     state["last_processed_file"] = None
     state["is_playing"] = False
@@ -745,6 +778,7 @@ def file_upload_stop_file_processing(state):
 
 def stop_file_processing_example(state):
     """Stops file processing when audio is cleared."""
+    ensure_sessions_initialized()
     state["mode"] = "idle"
     state["last_processed_file"] = None
     print(f"[DEBUG] global RUDE_LOOP_STOP is set to True")
@@ -786,8 +820,10 @@ def hide_examples_panel():
     """Collapse the example picker after an example is selected."""
     return False, gr.update(visible=False), "Choose Example"
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def reset_on_tab_switch(state):
     """Stop active processing and recreate both streaming instances."""
+    ensure_sessions_initialized()
     global MICROPHONE_PROCESSING_ENABLED, RUDE_LOOP_STOP
     RUDE_LOOP_STOP = True
     MICROPHONE_PROCESSING_ENABLED = False
@@ -836,8 +872,10 @@ def trigger_end_of_file_flag(state):
     print(f"[DEBUG] END OF FILE TRIGGERED trigger_end_of_file_flag(): end_of_file set to {state['end_of_file']}")
     return state
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def start_stream_processing(state):
     """Initiates stream processing mode."""
+    ensure_sessions_initialized()
     global MICROPHONE_PROCESSING_ENABLED
     print(f"[DEBUG] start_stream_processing() CALLED - mode before: {state.get('mode', 'None')}")
     MICROPHONE_PROCESSING_ENABLED = True
@@ -853,8 +891,10 @@ def start_stream_processing(state):
     print(f"[DEBUG] start_stream_processing() RETURNING - mode after: '{state['mode']}'")
     return state, gr.Button(value="Pause", visible=True)
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def toggle_microphone_pause(state):
     """Pause or resume microphone transcription while recording continues."""
+    ensure_sessions_initialized()
     if state.get("mode") != "stream":
         return (
             state,
@@ -883,8 +923,10 @@ def toggle_microphone_pause(state):
         format_transcript(state.get("full_transcript", "")),
     )
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def stop_stream_processing(state):
     """Stops stream processing mode."""
+    ensure_sessions_initialized()
     global MICROPHONE_PROCESSING_ENABLED
     print(f"[DEBUG] stop_stream_processing() CALLED - mode before: {state.get('mode', 'None')}")
     MICROPHONE_PROCESSING_ENABLED = False
@@ -907,8 +949,10 @@ def stop_stream_processing(state):
 
     return state, diar_html_output, transcript_textbox_html, gr.Button(value="Pause", visible=False)
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def stream_processing(state, audio_data):
     """Handles microphone streaming."""
+    ensure_sessions_initialized()
     # Initialize state if it's None or doesn't have required fields
     if state is None:
         state = {"mode": "stream", "full_transcript": "", "diar_html_initial": ""}
@@ -1015,6 +1059,7 @@ def reload_models_with_checkbox_update(
     )
     return html_string, single_speaker_mode
 
+@spaces.GPU(duration=GPU_DURATION_SECONDS)
 def reload_models(
     audio_pre_loading,
     streaming_chunk_size,
@@ -1025,6 +1070,7 @@ def reload_models(
     sent_break_sec,
 ):
     """Reloads the streaming session with new configurations."""
+    ensure_sessions_initialized()
     try:
         num_speakers = min(int(num_speakers), MAX_NUM_OF_SPEAKERS)
         if single_speaker_mode:
@@ -1106,13 +1152,22 @@ def get_text_box_html(
 
 def load_and_reset():
     """Resets the UI state on page load."""
-    AUDIO_PRE_LOADING = session_file._audio_pre_loading
-    STREAMING_CHUNK_SIZE = session_file.multispk_asr_streamer._att_context_size[1]
-    SINGLE_SPEAKER_MODE = 'ON' if session_file.multispk_asr_streamer._single_speaker_mode else 'OFF'
-    NUMBER_OF_SPEAKERS = f"{session_file.multispk_asr_streamer._max_num_of_spks}"
-    TEXT_BREAK_SILENCE = f"{session_file.multispk_asr_streamer._sent_break_sec} seconds"
-    CACHE_GATING_STATUS = 'ON' if session_file.multispk_asr_streamer._cache_gating else 'OFF'
-    BINARY_DIAR_PREDICTION_STATUS = 'ON' if session_file.multispk_asr_streamer._binary_diar_preds else 'OFF'
+    if MODELS_INITIALIZED:
+        AUDIO_PRE_LOADING = session_file._audio_pre_loading
+        STREAMING_CHUNK_SIZE = session_file.multispk_asr_streamer._att_context_size[1]
+        SINGLE_SPEAKER_MODE = 'ON' if session_file.multispk_asr_streamer._single_speaker_mode else 'OFF'
+        NUMBER_OF_SPEAKERS = f"{session_file.multispk_asr_streamer._max_num_of_spks}"
+        TEXT_BREAK_SILENCE = f"{session_file.multispk_asr_streamer._sent_break_sec} seconds"
+        CACHE_GATING_STATUS = 'ON' if session_file.multispk_asr_streamer._cache_gating else 'OFF'
+        BINARY_DIAR_PREDICTION_STATUS = 'ON' if session_file.multispk_asr_streamer._binary_diar_preds else 'OFF'
+    else:
+        AUDIO_PRE_LOADING = config_file.audio_pre_loading
+        STREAMING_CHUNK_SIZE = config_file.att_context_size[1]
+        SINGLE_SPEAKER_MODE = 'ON' if config_file.single_speaker_mode else 'OFF'
+        NUMBER_OF_SPEAKERS = f"{config_file.max_num_of_spks}"
+        TEXT_BREAK_SILENCE = f"{config_file.sent_break_sec} seconds"
+        CACHE_GATING_STATUS = 'ON' if config_file.cache_gating else 'OFF'
+        BINARY_DIAR_PREDICTION_STATUS = 'ON' if config_file.binary_diar_preds else 'OFF'
 
     status_textbox_html= get_text_box_html(
         audio_pre_loading=AUDIO_PRE_LOADING,
@@ -1124,7 +1179,7 @@ def load_and_reset():
         text_break_silence=TEXT_BREAK_SILENCE,
         )
     diar_html_output = render_diarization_state(None, config_file)
-    transcript_textbox = "<p style='text-align: center; font-size:20px;'><b>No transcription yet.</b></p>"
+    transcript_textbox = "<p style='text-align: center; font-size:20px;'><b>No transcription yet. Models load on first transcription request.</b></p>"
     return {"should_stop": False,
             "stream_step_num": 0,
             "microphone_paused": False,