fix: implement generation progress syncing and MPS memory tracking

Author: mzeena

backend/app_handler.py

@@ -162,7 +162,10 @@ def __init__(
             device=config.device,
         )
 
-        self.generation = GenerationHandler(state=self.state, lock=self._lock)
+        from state.job_queue import JobQueue
+        self.job_queue = JobQueue(persistence_path=config.settings_file.parent / "job_queue.json")
+
+        self.generation = GenerationHandler(state=self.state, lock=self._lock, job_queue=self.job_queue)
 
         self.video_generation = VideoGenerationHandler(
             state=self.state,
@@ -243,9 +246,6 @@ def __init__(
 
         self.downloads.cleanup_downloading_dir()
 
-        from state.job_queue import JobQueue
-        self.job_queue = JobQueue(persistence_path=config.settings_file.parent / "job_queue.json")
-
         # Wire up the QueueWorker with concrete executors so submitted jobs
         # are dispatched to the appropriate generation handler.
         from handlers.job_executors import ApiJobExecutor, GpuJobExecutor

backend/handlers/generation_handler.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 from api_types import CancelResponse, GenerationProgressResponse
 from handlers.base import StateHandlerBase, with_state_lock
@@ -17,11 +17,23 @@
     GpuSlot,
 )
 
+if TYPE_CHECKING:
+    from threading import RLock
+    from state.app_state_types import AppState
+    from state.job_queue import JobQueue
+
 logger = logging.getLogger(__name__)
 GenerationSlot = Literal["gpu", "api"]
 
-
 class GenerationHandler(StateHandlerBase):
+    def __init__(self, state: AppState, lock: RLock, job_queue: JobQueue | None = None) -> None:
+        super().__init__(state, lock)
+        self._job_queue = job_queue
+        self._current_job_id: str | None = None
+
+    def set_current_job_id(self, job_id: str | None) -> None:
+        self._current_job_id = job_id or None
+
     @with_state_lock
     def start_generation(self, generation_id: str) -> None:
         if self.is_generation_running():
@@ -91,6 +103,14 @@ def update_progress(
         current_step: int | None = None,
         total_steps: int | None = None,
     ) -> None:
+        # Sync to the persistent JobQueue if we have a job context
+        if self._job_queue and self._current_job_id:
+            self._job_queue.update_job(
+                self._current_job_id,
+                progress=progress,
+                phase=phase,
+            )
+
         match self._running_slot():
             case "gpu":
                 match self.state.gpu_slot:

backend/handlers/job_executors.py

@@ -54,9 +54,15 @@ def __init__(
 
     def execute(self, job: QueueJob) -> list[str]:
         logger.info("[QueueWorker] Executing GPU job %s (type=%s model=%s)", job.id, job.type, job.model)
-        if job.type == "image":
-            return self._execute_image(job)
-        return self._execute_video(job)
+        # Pass the queue job ID to the generation handler so it can sync progress
+        gen_handler = self._video._generation
+        gen_handler.set_current_job_id(job.id)
+        try:
+            if job.type == "image":
+                return self._execute_image(job)
+            return self._execute_video(job)
+        finally:
+            gen_handler.set_current_job_id(None)
 
     def _execute_video(self, job: QueueJob) -> list[str]:
         p = job.params
@@ -104,35 +110,40 @@ def __init__(
         self._image = image_generation
 
     def execute(self, job: QueueJob) -> list[str]:
-        logger.info("[QueueWorker] Executing API job %s (type=%s model=%s)", job.id, job.type, job.model)
-        if job.type == "image":
-            p = job.params
-            req = GenerateImageRequest(
-                prompt=_str(p, "prompt"),
-                width=_int(p, "width", 1920),
-                height=_int(p, "height", 1080),
-                numImages=_int(p, "numImages", 1),
-                numSteps=_int(p, "numSteps", 4),
-            )
-            result = self._image.generate(req)
-            return list(result.image_paths or [])
-        else:
-            p = job.params
-            req = GenerateVideoRequest(
-                prompt=_str(p, "prompt"),
-                imagePath=_str(p, "imagePath") or None,
-                lastFramePath=_str(p, "lastFramePath") or None,
-                audioPath=_str(p, "audioPath") or None,
-                resolution=_str(p, "resolution", "540p"),
-                duration=_str(p, "duration", "5"),
-                fps=_str(p, "fps", "24"),
-                audio=_str(p, "audio", "false"),
-                cameraMotion=_camera_motion(p),
-                aspectRatio=_aspect_ratio(p),
-                model=job.model,
-                negativePrompt=_str(p, "negativePrompt"),
-            )
-            result = self._video.generate(req)
-            if result.video_path:
-                return [result.video_path]
-            return []
+        logger.info("[QueueWorker] Executing API job %s (type=%s model=%s)", job.id, job.id, job.type, job.model)
+        gen_handler = self._video._generation
+        gen_handler.set_current_job_id(job.id)
+        try:
+            if job.type == "image":
+                p = job.params
+                req = GenerateImageRequest(
+                    prompt=_str(p, "prompt"),
+                    width=_int(p, "width", 1920),
+                    height=_int(p, "height", 1080),
+                    numImages=_int(p, "numImages", 1),
+                    numSteps=_int(p, "numSteps", 4),
+                )
+                result = self._image.generate(req)
+                return list(result.image_paths or [])
+            else:
+                p = job.params
+                req = GenerateVideoRequest(
+                    prompt=_str(p, "prompt"),
+                    imagePath=_str(p, "imagePath") or None,
+                    lastFramePath=_str(p, "lastFramePath") or None,
+                    audioPath=_str(p, "audioPath") or None,
+                    resolution=_str(p, "resolution", "540p"),
+                    duration=_str(p, "duration", "5"),
+                    fps=_str(p, "fps", "24"),
+                    audio=_str(p, "audio", "false"),
+                    cameraMotion=_camera_motion(p),
+                    aspectRatio=_aspect_ratio(p),
+                    model=job.model,
+                    negativePrompt=_str(p, "negativePrompt"),
+                )
+                result = self._video.generate(req)
+                if result.video_path:
+                    return [result.video_path]
+                return []
+        finally:
+            gen_handler.set_current_job_id(None)

backend/services/gpu_info/gpu_info_impl.py

@@ -79,10 +79,16 @@ def get_gpu_info(self) -> GpuTelemetryPayload:
         if self.get_mps_available():
             chip = self._get_macos_chip_name()
             name = f"{chip} (MPS)" if chip else "Apple Silicon (MPS)"
+            vram_used = 0
+            try:
+                # torch.mps.current_allocated_memory() provides allocated memory in bytes
+                vram_used = torch.mps.current_allocated_memory() // (1024 * 1024)
+            except Exception:
+                pass
             return {
                 "name": name,
                 "vram": self._get_system_ram_mb(),
-                "vramUsed": 0,
+                "vramUsed": vram_used,
             }
 
         return {"name": "Unknown", "vram": 0, "vramUsed": 0}