Upload app.py

app.py

@@ -56,8 +56,13 @@ LLM_REPO_ID = os.environ.get("LLM_REPO_ID", "aijadugar/wisprflow-clone-llm")
 # Under ZeroGPU, torch.cuda.is_available() is unreliable/False in the main
 # process anyway; actual device placement happens lazily inside the
 # @spaces.GPU-decorated functions below.
+#
+# bfloat16 (not float32) on CPU: halves RAM for both models. Modern PyTorch
+# has native CPU kernels for bf16 matmul/linear (unlike float16, which is
+# poorly supported on CPU), so this is safe and meaningfully cuts memory
+# pressure — important since the LLM checkpoint here is ~15GB.
 LOAD_DEVICE = "cpu"
-LOAD_DTYPE = torch.float32
+LOAD_DTYPE = torch.bfloat16
 
 MODE_PROMPTS = {
     "Email": (
@@ -98,6 +103,7 @@ whisper_processor = WhisperProcessor.from_pretrained(WHISPER_REPO_ID)
 whisper_model = WhisperForConditionalGeneration.from_pretrained(
     WHISPER_REPO_ID,
     torch_dtype=LOAD_DTYPE,
+    low_cpu_mem_usage=True,
 )
 whisper_model.eval()
 
@@ -106,6 +112,7 @@ llm_tokenizer = AutoTokenizer.from_pretrained(LLM_REPO_ID)
 llm_model = AutoModelForCausalLM.from_pretrained(
     LLM_REPO_ID,
     torch_dtype=LOAD_DTYPE,
+    low_cpu_mem_usage=True,
 )
 llm_model.eval()
 
@@ -191,31 +198,41 @@ def _clean_up_on_gpu(raw_text, mode):
 
 def transcribe(audio_path):
     if audio_path is None:
-        return "", 0.0
+        return "", 0.0, None
     start = time.time()
-    text = _transcribe_on_gpu(audio_path)
+    try:
+        text = _transcribe_on_gpu(audio_path)
+    except Exception as e:
+        return "", time.time() - start, f"Transcription failed: {e}"
     elapsed = time.time() - start
-    return text, elapsed
+    return text, elapsed, None
 
 
 def clean_up(raw_text, mode):
     if not raw_text.strip():
-        return "", 0.0
+        return "", 0.0, None
     start = time.time()
-    cleaned = _clean_up_on_gpu(raw_text, mode)
+    try:
+        cleaned = _clean_up_on_gpu(raw_text, mode)
+    except Exception as e:
+        return "", time.time() - start, f"Cleanup failed: {e}"
     elapsed = time.time() - start
-    return cleaned, elapsed
+    return cleaned, elapsed, None
 
 
 def run_pipeline(audio_path, mode):
     if audio_path is None:
         return "", "", "Record or upload audio first."
 
-    raw_text, asr_seconds = transcribe(audio_path)
+    raw_text, asr_seconds, asr_error = transcribe(audio_path)
+    if asr_error:
+        return "", "", asr_error
     if not raw_text:
         return "", "", "Couldn't transcribe that clip — try again."
 
-    cleaned_text, llm_seconds = clean_up(raw_text, mode)
+    cleaned_text, llm_seconds, llm_error = clean_up(raw_text, mode)
+    if llm_error:
+        return raw_text, "", llm_error
 
     total = asr_seconds + llm_seconds
     stats = (