feat: CUDA backend swap via binary download and restart

Add the ability to download a CUDA-enabled backend binary (~2.4 GB) and swap it in via a backend-only restart, solving the #1 user pain point (19 open 'GPU not detected' issues caused by GitHub's 2 GB asset limit). Backend: - cuda_download.py: download from R2 (primary) or GitHub split-parts (fallback), SHA-256 verification, atomic writes, progress via SSE - 4 new endpoints: GET/POST/DELETE /backend/cuda-*, GET cuda-progress - server.py: --version flag, auto-detect variant from binary name - build_binary.py: --cuda flag for CUDA PyInstaller builds - split_binary.py: split large binaries into <2GB GitHub Release assets - CI workflow for building CUDA binary Tauri: - restart_server command (stop -> wait -> start) - start_server prefers CUDA binary from {data_dir}/backends/ if present - Version mismatch check: runs --version before launching CUDA binary Frontend: - GpuAcceleration component: download, progress, restart, switch, delete - API client + types for CUDA status and management - Platform lifecycle: restartServer() on Tauri/Web - Aggressive 1s health polling during restart for fast reconnection
2026-09-20 07:10:40 -07:00 · 2026-03-13 00:04:12 -07:00
parent 758577fd4b
commit 2867421550
20 changed files with 2751 additions and 56 deletions
@@ -0,0 +1,387 @@
+import { useQuery, useQueryClient } from '@tanstack/react-query';
+import { AlertCircle, Cpu, Download, Loader2, RotateCw, Trash2, Zap } from 'lucide-react';
+import { useCallback, useEffect, useRef, useState } from 'react';
+import { Badge } from '@/components/ui/badge';
+import { Button } from '@/components/ui/button';
+import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
+import { Progress } from '@/components/ui/progress';
+import { apiClient } from '@/lib/api/client';
+import type { CudaDownloadProgress } from '@/lib/api/types';
+import { useServerHealth } from '@/lib/hooks/useServer';
+import { usePlatform } from '@/platform/PlatformContext';
+import { useServerStore } from '@/stores/serverStore';
+
+type RestartPhase = 'idle' | 'stopping' | 'waiting' | 'ready';
+
+export function GpuAcceleration() {
+  const platform = usePlatform();
+  const queryClient = useQueryClient();
+  const serverUrl = useServerStore((state) => state.serverUrl);
+  const { data: health } = useServerHealth();
+
+  const [restartPhase, setRestartPhase] = useState<RestartPhase>('idle');
+  const [error, setError] = useState<string | null>(null);
+  const [downloadProgress, setDownloadProgress] = useState<CudaDownloadProgress | null>(null);
+  const healthPollRef = useRef<ReturnType<typeof setInterval> | null>(null);
+
+  // Query CUDA backend status
+  const {
+    data: cudaStatus,
+    isLoading: cudaStatusLoading,
+    refetch: refetchCudaStatus,
+  } = useQuery({
+    queryKey: ['cuda-status', serverUrl],
+    queryFn: () => apiClient.getCudaStatus(),
+    refetchInterval: cudaStatusLoading ? false : 10000,
+    retry: 1,
+    enabled: !!health, // Only fetch when backend is reachable
+  });
+
+  // Derived state
+  const isCurrentlyCuda = health?.backend_variant === 'cuda';
+  const cudaAvailable = cudaStatus?.available ?? false;
+  const cudaDownloading = cudaStatus?.downloading ?? false;
+
+  // Clean up health poll on unmount
+  useEffect(() => {
+    return () => {
+      if (healthPollRef.current) {
+        clearInterval(healthPollRef.current);
+        healthPollRef.current = null;
+      }
+    };
+  }, []);
+
+  // SSE progress tracking during download
+  useEffect(() => {
+    if (!cudaDownloading || !serverUrl) {
+      return;
+    }
+
+    const eventSource = new EventSource(`${serverUrl}/backend/cuda-progress`);
+
+    eventSource.onmessage = (event) => {
+      try {
+        const data = JSON.parse(event.data) as CudaDownloadProgress;
+        setDownloadProgress(data);
+
+        if (data.status === 'complete') {
+          eventSource.close();
+          setDownloadProgress(null);
+          refetchCudaStatus();
+        } else if (data.status === 'error') {
+          eventSource.close();
+          setError(data.error || 'Download failed');
+          setDownloadProgress(null);
+          refetchCudaStatus();
+        }
+      } catch (e) {
+        console.error('Error parsing CUDA progress event:', e);
+      }
+    };
+
+    eventSource.onerror = () => {
+      eventSource.close();
+    };
+
+    return () => {
+      eventSource.close();
+    };
+  }, [cudaDownloading, serverUrl, refetchCudaStatus]);
+
+  // Start aggressive health polling during restart
+  const startHealthPolling = useCallback(() => {
+    if (healthPollRef.current) return;
+
+    healthPollRef.current = setInterval(async () => {
+      try {
+        const result = await apiClient.getHealth();
+        if (result.status === 'healthy') {
+          // Server is back up
+          if (healthPollRef.current) {
+            clearInterval(healthPollRef.current);
+            healthPollRef.current = null;
+          }
+          setRestartPhase('ready');
+          // Invalidate all queries to refresh UI
+          queryClient.invalidateQueries();
+          // Reset after a moment
+          setTimeout(() => setRestartPhase('idle'), 2000);
+        }
+      } catch {
+        // Server still down, keep polling
+      }
+    }, 1000);
+  }, [queryClient]);
+
+  const handleDownload = async () => {
+    setError(null);
+    try {
+      await apiClient.downloadCudaBackend();
+      refetchCudaStatus();
+    } catch (e: unknown) {
+      const msg = e instanceof Error ? e.message : 'Failed to start download';
+      if (msg.includes('already downloaded')) {
+        refetchCudaStatus();
+      } else {
+        setError(msg);
+      }
+    }
+  };
+
+  const handleRestart = async () => {
+    setError(null);
+    setRestartPhase('stopping');
+
+    try {
+      setRestartPhase('waiting');
+      startHealthPolling();
+      await platform.lifecycle.restartServer();
+      // Invoke resolved — server is likely ready. Stop polling and refresh.
+      if (healthPollRef.current) {
+        clearInterval(healthPollRef.current);
+        healthPollRef.current = null;
+      }
+      setRestartPhase('ready');
+      queryClient.invalidateQueries();
+      setTimeout(() => setRestartPhase('idle'), 2000);
+    } catch (e: unknown) {
+      setRestartPhase('idle');
+      if (healthPollRef.current) {
+        clearInterval(healthPollRef.current);
+        healthPollRef.current = null;
+      }
+      setError(e instanceof Error ? e.message : 'Restart failed');
+    }
+  };
+
+  const handleSwitchToCpu = async () => {
+    // To switch to CPU: delete the CUDA binary, then restart.
+    // start_server always prefers CUDA if present, so we must remove it first.
+    setError(null);
+    setRestartPhase('stopping');
+
+    try {
+      await apiClient.deleteCudaBackend();
+      setRestartPhase('waiting');
+      startHealthPolling();
+      await platform.lifecycle.restartServer();
+      // Invoke resolved — server is likely ready
+      if (healthPollRef.current) {
+        clearInterval(healthPollRef.current);
+        healthPollRef.current = null;
+      }
+      setRestartPhase('ready');
+      queryClient.invalidateQueries();
+      setTimeout(() => setRestartPhase('idle'), 2000);
+    } catch (e: unknown) {
+      setRestartPhase('idle');
+      if (healthPollRef.current) {
+        clearInterval(healthPollRef.current);
+        healthPollRef.current = null;
+      }
+      setError(e instanceof Error ? e.message : 'Failed to switch to CPU');
+      refetchCudaStatus();
+    }
+  };
+
+  const handleDelete = async () => {
+    setError(null);
+    try {
+      await apiClient.deleteCudaBackend();
+      refetchCudaStatus();
+    } catch (e: unknown) {
+      setError(e instanceof Error ? e.message : 'Failed to delete CUDA backend');
+    }
+  };
+
+  const formatBytes = (bytes: number): string => {
+    if (bytes === 0) return '0 B';
+    const k = 1024;
+    const sizes = ['B', 'KB', 'MB', 'GB'];
+    const i = Math.floor(Math.log(bytes) / Math.log(k));
+    return `${(bytes / k ** i).toFixed(1)} ${sizes[i]}`;
+  };
+
+  // Don't render until health data is available
+  if (!health) return null;
+
+  // If the system already has native GPU (MPS, etc.), only show info - no CUDA needed
+  const hasNativeGpu =
+    health.gpu_available &&
+    !isCurrentlyCuda &&
+    health.gpu_type &&
+    !health.gpu_type.includes('CUDA');
+
+  return (
+    <Card>
+      <CardHeader>
+        <CardTitle className="flex items-center gap-2">
+          <Zap className="h-4 w-4" />
+          GPU Acceleration
+        </CardTitle>
+      </CardHeader>
+      <CardContent className="space-y-4">
+        {/* Current status */}
+        <div className="flex items-center justify-between">
+          <div className="space-y-1">
+            <div className="text-sm font-medium">Backend</div>
+            <div className="text-sm text-muted-foreground">
+              {isCurrentlyCuda ? 'CUDA (GPU accelerated)' : 'CPU'}
+            </div>
+          </div>
+          <Badge variant={isCurrentlyCuda ? 'default' : 'secondary'}>
+            {isCurrentlyCuda ? (
+              <>
+                <Zap className="h-3 w-3 mr-1" /> CUDA
+              </>
+            ) : (
+              <>
+                <Cpu className="h-3 w-3 mr-1" /> CPU
+              </>
+            )}
+          </Badge>
+        </div>
+
+        {/* GPU info from health */}
+        {health.gpu_type && (
+          <div className="space-y-1">
+            <div className="text-sm font-medium">GPU</div>
+            <div className="text-sm text-muted-foreground">{health.gpu_type}</div>
+            {health.vram_used_mb != null && (
+              <div className="text-xs text-muted-foreground">
+                VRAM: {health.vram_used_mb.toFixed(0)} MB used
+              </div>
+            )}
+          </div>
+        )}
+
+        {/* Native GPU detected - no CUDA download needed */}
+        {hasNativeGpu && (
+          <div className="p-3 rounded-lg bg-accent/10 border border-accent/20">
+            <div className="text-sm">
+              Your system uses <strong>{health.gpu_type}</strong> for acceleration. No additional
+              downloads needed.
+            </div>
+          </div>
+        )}
+
+        {/* CUDA download section - only show when native GPU is NOT detected (i.e., Windows/Linux NVIDIA users) */}
+        {!hasNativeGpu && (
+          <>
+            {/* Download progress */}
+            {cudaDownloading && downloadProgress && (
+              <div className="space-y-2">
+                <div className="flex items-center justify-between text-sm">
+                  <div className="flex items-center gap-2">
+                    <Loader2 className="h-4 w-4 animate-spin" />
+                    <span>{downloadProgress.filename || 'Downloading CUDA backend...'}</span>
+                  </div>
+                  {downloadProgress.total > 0 && (
+                    <span className="text-muted-foreground">
+                      {downloadProgress.progress.toFixed(1)}%
+                    </span>
+                  )}
+                </div>
+                {downloadProgress.total > 0 && (
+                  <>
+                    <Progress value={downloadProgress.progress} className="h-2" />
+                    <div className="text-xs text-muted-foreground">
+                      {formatBytes(downloadProgress.current)} /{' '}
+                      {formatBytes(downloadProgress.total)}
+                    </div>
+                  </>
+                )}
+              </div>
+            )}
+
+            {/* Restart in progress */}
+            {restartPhase !== 'idle' && (
+              <div className="flex items-center gap-2 p-3 rounded-lg bg-primary/5 border">
+                <Loader2 className="h-4 w-4 animate-spin" />
+                <span className="text-sm">
+                  {restartPhase === 'stopping' && 'Stopping server...'}
+                  {restartPhase === 'waiting' && 'Restarting server...'}
+                  {restartPhase === 'ready' && 'Server restarted successfully!'}
+                </span>
+              </div>
+            )}
+
+            {/* Error display */}
+            {error && (
+              <div className="flex items-center gap-2 text-sm text-destructive">
+                <AlertCircle className="h-4 w-4 shrink-0" />
+                <span>{error}</span>
+              </div>
+            )}
+
+            {/* Actions */}
+            {restartPhase === 'idle' && !cudaDownloading && (
+              <div className="space-y-2">
+                {/* Not downloaded yet - show download button */}
+                {!cudaAvailable && (
+                  <div className="space-y-3">
+                    <p className="text-sm text-muted-foreground">
+                      Download the CUDA backend (~2.4 GB) for NVIDIA GPU acceleration. Requires an
+                      NVIDIA GPU with CUDA support.
+                    </p>
+                    <Button onClick={handleDownload} className="w-full" size="sm">
+                      <Download className="h-4 w-4 mr-2" />
+                      Download CUDA Backend
+                    </Button>
+                  </div>
+                )}
+
+                {/* Downloaded but not active - show switch button */}
+                {cudaAvailable && !isCurrentlyCuda && platform.metadata.isTauri && (
+                  <div className="space-y-3">
+                    <p className="text-sm text-muted-foreground">
+                      CUDA backend is downloaded and ready. Restart the server to enable GPU
+                      acceleration.
+                    </p>
+                    <Button onClick={handleRestart} className="w-full" size="sm">
+                      <RotateCw className="h-4 w-4 mr-2" />
+                      Switch to CUDA Backend
+                    </Button>
+                  </div>
+                )}
+
+                {/* Currently active - show switch back to CPU */}
+                {isCurrentlyCuda && platform.metadata.isTauri && (
+                  <div className="space-y-3">
+                    <p className="text-sm text-muted-foreground">
+                      Running with CUDA GPU acceleration. Switch back to CPU if needed (you can
+                      re-download later).
+                    </p>
+                    <Button
+                      onClick={handleSwitchToCpu}
+                      variant="outline"
+                      className="w-full"
+                      size="sm"
+                    >
+                      <RotateCw className="h-4 w-4 mr-2" />
+                      Switch to CPU Backend
+                    </Button>
+                  </div>
+                )}
+
+                {/* Delete option when downloaded (and not active) */}
+                {cudaAvailable && !isCurrentlyCuda && (
+                  <Button
+                    onClick={handleDelete}
+                    variant="ghost"
+                    className="w-full text-muted-foreground hover:text-destructive"
+                    size="sm"
+                  >
+                    <Trash2 className="h-4 w-4 mr-2" />
+                    Remove CUDA Backend
+                  </Button>
+                )}
+              </div>
+            )}
+          </>
+        )}
+      </CardContent>
+    </Card>
+  );
+}
@@ -1,4 +1,5 @@
 import { ConnectionForm } from '@/components/ServerSettings/ConnectionForm';
+import { GpuAcceleration } from '@/components/ServerSettings/GpuAcceleration';
 import { ServerStatus } from '@/components/ServerSettings/ServerStatus';
 import { UpdateStatus } from '@/components/ServerSettings/UpdateStatus';
 import { usePlatform } from '@/platform/PlatformContext';
@@ -11,6 +12,7 @@ export function ServerTab() {
        <ConnectionForm />
        <ServerStatus />
      </div>
+      {platform.metadata.isTauri && <GpuAcceleration />}
      {platform.metadata.isTauri && <UpdateStatus />}
      <div className="py-8 text-center text-sm text-muted-foreground">
        Created by{' '}
@@ -1,29 +1,30 @@
-import { useServerStore } from '@/stores/serverStore';
 import type { LanguageCode } from '@/lib/constants/languages';
+import { useServerStore } from '@/stores/serverStore';
 import type {
-  VoiceProfileCreate,
-  VoiceProfileResponse,
-  ProfileSampleResponse,
+  ActiveTasksResponse,
+  CudaStatus,
  GenerationRequest,
  GenerationResponse,
-  HistoryQuery,
-  HistoryListResponse,
-  HistoryResponse,
-  TranscriptionResponse,
  HealthResponse,
-  ModelStatusListResponse,
+  HistoryListResponse,
+  HistoryQuery,
+  HistoryResponse,
  ModelDownloadRequest,
-  ActiveTasksResponse,
+  ModelStatusListResponse,
+  ProfileSampleResponse,
  StoryCreate,
-  StoryResponse,
  StoryDetailResponse,
+  StoryItemBatchUpdate,
  StoryItemCreate,
  StoryItemDetail,
-  StoryItemBatchUpdate,
-  StoryItemReorder,
  StoryItemMove,
-  StoryItemTrim,
+  StoryItemReorder,
  StoryItemSplit,
+  StoryItemTrim,
+  StoryResponse,
+  TranscriptionResponse,
+  VoiceProfileCreate,
+  VoiceProfileResponse,
 } from './types';

 class ApiClient {
@@ -251,7 +252,13 @@ class ApiClient {
    return response.blob();
  }

-  async importGeneration(file: File): Promise<{ id: string; profile_id: string; profile_name: string; text: string; message: string }> {
+  async importGeneration(file: File): Promise<{
+    id: string;
+    profile_id: string;
+    profile_name: string;
+    text: string;
+    message: string;
+  }> {
    const url = `${this.getBaseUrl()}/history/import`;
    const formData = new FormData();
    formData.append('file', file);
@@ -310,7 +317,12 @@ class ApiClient {
  }

  async triggerModelDownload(modelName: string): Promise<{ message: string }> {
-    console.log('[API] triggerModelDownload called for:', modelName, 'at', new Date().toISOString());
+    console.log(
+      '[API] triggerModelDownload called for:',
+      modelName,
+      'at',
+      new Date().toISOString(),
+    );
    const result = await this.request<{ message: string }>('/models/download', {
      method: 'POST',
      body: JSON.stringify({ model_name: modelName } as ModelDownloadRequest),
@@ -354,10 +366,7 @@ class ApiClient {
    return this.request('/channels');
  }

-  async createChannel(data: {
-    name: string;
-    device_ids: string[];
-  }): Promise<{
+  async createChannel(data: { name: string; device_ids: string[] }): Promise<{
    id: string;
    name: string;
    is_default: boolean;
@@ -399,10 +408,7 @@ class ApiClient {
    return this.request(`/channels/${channelId}/voices`);
  }

-  async setChannelVoices(
-    channelId: string,
-    profileIds: string[],
-  ): Promise<{ message: string }> {
+  async setChannelVoices(channelId: string, profileIds: string[]): Promise<{ message: string }> {
    return this.request(`/channels/${channelId}/voices`, {
      method: 'PUT',
      body: JSON.stringify({ profile_ids: profileIds }),
@@ -413,16 +419,30 @@ class ApiClient {
    return this.request(`/profiles/${profileId}/channels`);
  }

-  async setProfileChannels(
-    profileId: string,
-    channelIds: string[],
-  ): Promise<{ message: string }> {
+  async setProfileChannels(profileId: string, channelIds: string[]): Promise<{ message: string }> {
    return this.request(`/profiles/${profileId}/channels`, {
      method: 'PUT',
      body: JSON.stringify({ channel_ids: channelIds }),
    });
  }

+  // CUDA Backend Management
+  async getCudaStatus(): Promise<CudaStatus> {
+    return this.request<CudaStatus>('/backend/cuda-status');
+  }
+
+  async downloadCudaBackend(): Promise<{ message: string; progress_key: string }> {
+    return this.request<{ message: string; progress_key: string }>('/backend/download-cuda', {
+      method: 'POST',
+    });
+  }
+
+  async deleteCudaBackend(): Promise<{ message: string }> {
+    return this.request<{ message: string }>('/backend/cuda', {
+      method: 'DELETE',
+    });
+  }
+
  // Stories
  async listStories(): Promise<StoryResponse[]> {
    return this.request<StoryResponse[]>('/stories');
@@ -479,21 +499,33 @@ class ApiClient {
    });
  }

-  async moveStoryItem(storyId: string, itemId: string, data: StoryItemMove): Promise<StoryItemDetail> {
+  async moveStoryItem(
+    storyId: string,
+    itemId: string,
+    data: StoryItemMove,
+  ): Promise<StoryItemDetail> {
    return this.request<StoryItemDetail>(`/stories/${storyId}/items/${itemId}/move`, {
      method: 'PUT',
      body: JSON.stringify(data),
    });
  }

-  async trimStoryItem(storyId: string, itemId: string, data: StoryItemTrim): Promise<StoryItemDetail> {
+  async trimStoryItem(
+    storyId: string,
+    itemId: string,
+    data: StoryItemTrim,
+  ): Promise<StoryItemDetail> {
    return this.request<StoryItemDetail>(`/stories/${storyId}/items/${itemId}/trim`, {
      method: 'PUT',
      body: JSON.stringify(data),
    });
  }

-  async splitStoryItem(storyId: string, itemId: string, data: StoryItemSplit): Promise<StoryItemDetail[]> {
+  async splitStoryItem(
+    storyId: string,
+    itemId: string,
+    data: StoryItemSplit,
+  ): Promise<StoryItemDetail[]> {
    return this.request<StoryItemDetail[]>(`/stories/${storyId}/items/${itemId}/split`, {
      method: 'POST',
      body: JSON.stringify(data),
@@ -78,7 +78,29 @@ export interface HealthResponse {
  model_downloaded?: boolean;
  model_size?: string;
  gpu_available: boolean;
+  gpu_type?: string;
  vram_used_mb?: number;
+  backend_type?: string;
+  backend_variant?: string; // "cpu" or "cuda"
+}
+
+export interface CudaDownloadProgress {
+  model_name: string;
+  current: number;
+  total: number;
+  progress: number;
+  filename?: string;
+  status: 'downloading' | 'extracting' | 'complete' | 'error';
+  timestamp: string;
+  error?: string;
+}
+
+export interface CudaStatus {
+  available: boolean; // CUDA binary exists on disk
+  active: boolean; // Currently running the CUDA binary
+  binary_path?: string;
+  downloading: boolean; // Download in progress
+  download_progress?: CudaDownloadProgress;
 }

 export interface ModelProgress {
@@ -96,7 +118,7 @@ export interface ModelStatus {
  model_name: string;
  display_name: string;
  downloaded: boolean;
-  downloading: boolean;  // True if download is in progress
+  downloading: boolean; // True if download is in progress
  size_mb?: number;
  loaded: boolean;
 }
@@ -51,6 +51,7 @@ export interface PlatformAudio {
 export interface PlatformLifecycle {
  startServer(remote?: boolean): Promise<string>;
  stopServer(): Promise<void>;
+  restartServer(): Promise<string>;
  setKeepServerRunning(keep: boolean): Promise<void>;
  setupWindowCloseHandler(): Promise<void>;
  onServerReady?: () => void;