From cf3cf3f0022576b58cfeaccbf260feb3ec5e936f Mon Sep 17 00:00:00 2001
From: Jamie Pine <ijamespine@me.com>
Date: Thu, 29 Jan 2026 15:57:53 -0800
Subject: [PATCH] Enhance model download handling in useGenerationForm and
 ProgressManager

- Introduced a flag to track download status in useGenerationForm, ensuring proper UI updates during model downloads.
- Updated ProgressManager to only send initial progress updates if the model is actively downloading or extracting, preventing outdated status messages from being sent.
- Improved error handling and logging for better visibility into model download processes.
---
 app/src/lib/hooks/useGenerationForm.ts | 11 +++++++++--
 backend/utils/progress.py              | 12 +++++++++---
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/app/src/lib/hooks/useGenerationForm.ts b/app/src/lib/hooks/useGenerationForm.ts
index c6fdba50..2126d90d 100644
--- a/app/src/lib/hooks/useGenerationForm.ts
+++ b/app/src/lib/hooks/useGenerationForm.ts
@@ -70,6 +70,7 @@ export function useGenerationForm(options: UseGenerationFormOptions = {}) {
       const modelName = `qwen-tts-${data.modelSize}`;
       const displayName = data.modelSize === '1.7B' ? 'Qwen TTS 1.7B' : 'Qwen TTS 0.6B';
 
+      let isDownloading = false;
       try {
         const modelStatus = await apiClient.getModelStatus();
         const model = modelStatus.models.find((m) => m.model_name === modelName);
@@ -77,6 +78,7 @@ export function useGenerationForm(options: UseGenerationFormOptions = {}) {
         if (model && !model.downloaded) {
           setDownloadingModelName(modelName);
           setDownloadingDisplayName(displayName);
+          isDownloading = true;
         }
       } catch (error) {
         console.error('Failed to check model status:', error);
@@ -101,16 +103,21 @@ export function useGenerationForm(options: UseGenerationFormOptions = {}) {
 
       form.reset();
       options.onSuccess?.(result.id);
+
+      if (isDownloading) {
+        setDownloadingModelName(null);
+        setDownloadingDisplayName(null);
+      }
     } catch (error) {
       toast({
         title: 'Generation failed',
         description: error instanceof Error ? error.message : 'Failed to generate audio',
         variant: 'destructive',
       });
-    } finally {
-      setIsGenerating(false);
       setDownloadingModelName(null);
       setDownloadingDisplayName(null);
+    } finally {
+      setIsGenerating(false);
     }
   }
 
diff --git a/backend/utils/progress.py b/backend/utils/progress.py
index bd51d984..dec1d1cc 100644
--- a/backend/utils/progress.py
+++ b/backend/utils/progress.py
@@ -121,10 +121,16 @@ class ProgressManager:
         logger.info(f"SSE client subscribed to {model_name}, total listeners: {len(self._listeners[model_name])}")
 
         try:
-            # Send initial progress if available
+            # Send initial progress if available and still in progress
             if model_name in self._progress:
-                logger.info(f"Sending initial progress for {model_name}: {self._progress[model_name].get('status')}")
-                yield f"data: {json.dumps(self._progress[model_name])}\n\n"
+                status = self._progress[model_name].get('status')
+                # Only send initial progress if download is actually in progress
+                # Don't send old 'complete' or 'error' status from previous downloads
+                if status in ('downloading', 'extracting'):
+                    logger.info(f"Sending initial progress for {model_name}: {status}")
+                    yield f"data: {json.dumps(self._progress[model_name])}\n\n"
+                else:
+                    logger.info(f"Skipping initial progress for {model_name} (status: {status})")
             else:
                 logger.info(f"No initial progress available for {model_name}")