Implement sidebar navigation and model management features. Refactor App component to utilize a Sidebar for tab navigation, integrating ProfileList, GenerationForm, HistoryTable, and ServerStatus components. Introduce ModelManagement and ModelProgress components for handling AI model downloads and status updates. Enhance CSS for sidebar styling and add progress tracking functionality in the backend for model downloads.

2026-09-18 14:20:42 -07:00 · 2026-01-25 03:10:16 -08:00
parent ca3409ebef
commit 6429cb6673
12 changed files with 1061 additions and 82 deletions
@@ -0,0 +1,179 @@
+import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query';
+import { apiClient } from '@/lib/api/client';
+import { Card, CardContent, CardHeader, CardTitle, CardDescription } from '@/components/ui/card';
+import { Button } from '@/components/ui/button';
+import { Badge } from '@/components/ui/badge';
+import { Loader2, Download, CheckCircle2 } from 'lucide-react';
+import { ModelProgress } from './ModelProgress';
+import { useToast } from '@/components/ui/use-toast';
+
+export function ModelManagement() {
+  const { toast } = useToast();
+  const queryClient = useQueryClient();
+
+  const { data: modelStatus, isLoading } = useQuery({
+    queryKey: ['modelStatus'],
+    queryFn: () => apiClient.getModelStatus(),
+    refetchInterval: 5000, // Refresh every 5 seconds
+  });
+
+  const downloadMutation = useMutation({
+    mutationFn: (modelName: string) => apiClient.triggerModelDownload(modelName),
+    onSuccess: (_, modelName) => {
+      toast({
+        title: 'Download started',
+        description: `Downloading ${modelName}...`,
+      });
+      // Refetch status after a delay to see progress
+      setTimeout(() => {
+        queryClient.invalidateQueries({ queryKey: ['modelStatus'] });
+      }, 1000);
+    },
+    onError: (error: Error) => {
+      toast({
+        title: 'Download failed',
+        description: error.message,
+        variant: 'destructive',
+      });
+    },
+  });
+
+  const formatSize = (sizeMb?: number): string => {
+    if (!sizeMb) return 'Unknown';
+    if (sizeMb < 1024) return `${sizeMb.toFixed(1)} MB`;
+    return `${(sizeMb / 1024).toFixed(2)} GB`;
+  };
+
+  return (
+    <Card>
+      <CardHeader>
+        <CardTitle>Model Management</CardTitle>
+        <CardDescription>
+          Download and manage AI models for voice generation and transcription
+        </CardDescription>
+      </CardHeader>
+      <CardContent className="space-y-4">
+        {isLoading ? (
+          <div className="flex items-center justify-center py-8">
+            <Loader2 className="h-6 w-6 animate-spin text-muted-foreground" />
+          </div>
+        ) : modelStatus ? (
+          <div className="space-y-4">
+            {/* TTS Models */}
+            <div>
+              <h3 className="text-sm font-semibold mb-3 text-muted-foreground">Voice Generation Models</h3>
+              <div className="space-y-2">
+                {modelStatus.models
+                  .filter((m) => m.model_name.startsWith('qwen-tts'))
+                  .map((model) => (
+                    <ModelItem
+                      key={model.model_name}
+                      model={model}
+                      onDownload={() => downloadMutation.mutate(model.model_name)}
+                      isDownloading={downloadMutation.isPending}
+                      formatSize={formatSize}
+                    />
+                  ))}
+              </div>
+            </div>
+
+            {/* Whisper Models */}
+            <div>
+              <h3 className="text-sm font-semibold mb-3 text-muted-foreground">Transcription Models</h3>
+              <div className="space-y-2">
+                {modelStatus.models
+                  .filter((m) => m.model_name.startsWith('whisper'))
+                  .map((model) => (
+                    <ModelItem
+                      key={model.model_name}
+                      model={model}
+                      onDownload={() => downloadMutation.mutate(model.model_name)}
+                      isDownloading={downloadMutation.isPending}
+                      formatSize={formatSize}
+                    />
+                  ))}
+              </div>
+            </div>
+
+            {/* Progress indicators */}
+            <div className="pt-4 border-t">
+              <h3 className="text-sm font-semibold mb-3 text-muted-foreground">Download Progress</h3>
+              <div className="space-y-2">
+                {modelStatus.models.map((model) => (
+                  <ModelProgress
+                    key={model.model_name}
+                    modelName={model.model_name}
+                    displayName={model.display_name}
+                  />
+                ))}
+              </div>
+            </div>
+          </div>
+        ) : null}
+      </CardContent>
+    </Card>
+  );
+}
+
+interface ModelItemProps {
+  model: {
+    model_name: string;
+    display_name: string;
+    downloaded: boolean;
+    size_mb?: number;
+    loaded: boolean;
+  };
+  onDownload: () => void;
+  isDownloading: boolean;
+  formatSize: (sizeMb?: number) => string;
+}
+
+function ModelItem({ model, onDownload, isDownloading, formatSize }: ModelItemProps) {
+  return (
+    <div className="flex items-center justify-between p-3 border rounded-lg">
+      <div className="flex-1">
+        <div className="flex items-center gap-2">
+          <span className="font-medium text-sm">{model.display_name}</span>
+          {model.loaded && (
+            <Badge variant="default" className="text-xs">Loaded</Badge>
+          )}
+          {model.downloaded && !model.loaded && (
+            <Badge variant="secondary" className="text-xs">Downloaded</Badge>
+          )}
+        </div>
+        {model.downloaded && model.size_mb && (
+          <div className="text-xs text-muted-foreground mt-1">
+            Size: {formatSize(model.size_mb)}
+          </div>
+        )}
+      </div>
+      <div className="flex items-center gap-2">
+        {model.downloaded ? (
+          <div className="flex items-center gap-1 text-sm text-muted-foreground">
+            <CheckCircle2 className="h-4 w-4 text-green-500" />
+            <span>Ready</span>
+          </div>
+        ) : (
+          <Button
+            size="sm"
+            onClick={onDownload}
+            disabled={isDownloading}
+            variant="outline"
+          >
+            {isDownloading ? (
+              <>
+                <Loader2 className="h-4 w-4 mr-2 animate-spin" />
+                Downloading...
+              </>
+            ) : (
+              <>
+                <Download className="h-4 w-4 mr-2" />
+                Download
+              </>
+            )}
+          </Button>
+        )}
+      </div>
+    </div>
+  );
+}
@@ -0,0 +1,121 @@
+import { useEffect, useState } from 'react';
+import { Progress } from '@/components/ui/progress';
+import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
+import { useServerStore } from '@/stores/serverStore';
+import type { ModelProgress as ModelProgressType } from '@/lib/api/types';
+import { Loader2, CheckCircle2, XCircle } from 'lucide-react';
+
+interface ModelProgressProps {
+  modelName: string;
+  displayName: string;
+}
+
+export function ModelProgress({ modelName, displayName }: ModelProgressProps) {
+  const [progress, setProgress] = useState<ModelProgressType | null>(null);
+  const [isSubscribed, setIsSubscribed] = useState(false);
+  const serverUrl = useServerStore((state) => state.serverUrl);
+
+  useEffect(() => {
+    if (!serverUrl || isSubscribed) return;
+
+    // Subscribe to progress updates via Server-Sent Events
+    const eventSource = new EventSource(`${serverUrl}/models/progress/${modelName}`);
+
+    eventSource.onmessage = (event) => {
+      try {
+        const data = JSON.parse(event.data) as ModelProgressType;
+        setProgress(data);
+        
+        // Close connection if complete or error
+        if (data.status === 'complete' || data.status === 'error') {
+          eventSource.close();
+          setIsSubscribed(false);
+        }
+      } catch (error) {
+        console.error('Error parsing progress event:', error);
+      }
+    };
+
+    eventSource.onerror = (error) => {
+      console.error('SSE error:', error);
+      eventSource.close();
+      setIsSubscribed(false);
+    };
+
+    setIsSubscribed(true);
+
+    return () => {
+      eventSource.close();
+      setIsSubscribed(false);
+    };
+  }, [serverUrl, modelName, isSubscribed]);
+
+  // Don't render if no progress or if complete/error and some time has passed
+  if (!progress || (progress.status === 'complete' && Date.now() - new Date(progress.timestamp).getTime() > 5000)) {
+    return null;
+  }
+
+  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 / Math.pow(k, i)).toFixed(1)} ${sizes[i]}`;
+  };
+
+  const getStatusIcon = () => {
+    switch (progress.status) {
+      case 'complete':
+        return <CheckCircle2 className="h-4 w-4 text-green-500" />;
+      case 'error':
+        return <XCircle className="h-4 w-4 text-destructive" />;
+      case 'downloading':
+      case 'extracting':
+        return <Loader2 className="h-4 w-4 animate-spin" />;
+      default:
+        return null;
+    }
+  };
+
+  const getStatusText = () => {
+    switch (progress.status) {
+      case 'complete':
+        return 'Download complete';
+      case 'error':
+        return `Error: ${progress.error || 'Unknown error'}`;
+      case 'downloading':
+        return progress.filename ? `Downloading ${progress.filename}...` : 'Downloading...';
+      case 'extracting':
+        return 'Extracting...';
+      default:
+        return 'Processing...';
+    }
+  };
+
+  return (
+    <Card className="mb-4">
+      <CardHeader className="pb-3">
+        <CardTitle className="text-sm font-medium flex items-center gap-2">
+          {getStatusIcon()}
+          {displayName}
+        </CardTitle>
+      </CardHeader>
+      <CardContent className="space-y-2">
+        <div className="space-y-1">
+          <div className="flex justify-between text-xs text-muted-foreground">
+            <span>{getStatusText()}</span>
+            {progress.total > 0 && (
+              <span>
+                {formatBytes(progress.current)} / {formatBytes(progress.total)} (
+                {progress.progress.toFixed(1)}%)
+              </span>
+            )}
+          </div>
+          {progress.total > 0 && (
+            <Progress value={progress.progress} className="h-2" />
+          )}
+        </div>
+      </CardContent>
+    </Card>
+  );
+}
@@ -3,6 +3,7 @@ import { Badge } from '@/components/ui/badge';
 import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
 import { useServerHealth } from '@/lib/hooks/useServer';
 import { useServerStore } from '@/stores/serverStore';
+import { ModelProgress } from './ModelProgress';

 export function ServerStatus() {
  const { data: health, isLoading, error } = useServerHealth();
@@ -19,6 +20,16 @@ export function ServerStatus() {
          <div className="font-mono text-sm">{serverUrl}</div>
        </div>

+        {/* Model download progress */}
+        <div className="space-y-2">
+          <ModelProgress modelName="qwen-tts-1.7B" displayName="Qwen TTS 1.7B" />
+          <ModelProgress modelName="qwen-tts-0.6B" displayName="Qwen TTS 0.6B" />
+          <ModelProgress modelName="whisper-base" displayName="Whisper Base" />
+          <ModelProgress modelName="whisper-small" displayName="Whisper Small" />
+          <ModelProgress modelName="whisper-medium" displayName="Whisper Medium" />
+          <ModelProgress modelName="whisper-large" displayName="Whisper Large" />
+        </div>
+
        {isLoading ? (
          <div className="flex items-center gap-2">
            <Loader2 className="h-4 w-4 animate-spin" />
@@ -35,10 +46,19 @@ export function ServerStatus() {
              <CheckCircle2 className="h-4 w-4 text-green-500" />
              <span className="text-sm">Connected</span>
            </div>
-            <div className="flex gap-2">
+            <div className="flex flex-wrap gap-2">
              <Badge variant={health.model_loaded ? 'default' : 'secondary'}>
-                Model: {health.model_loaded ? 'Loaded' : 'Not Loaded'}
+                Model: {health.model_loaded 
+                  ? `Loaded${health.model_size ? ` (${health.model_size})` : ''}` 
+                  : health.model_downloaded === false 
+                    ? 'Not Downloaded' 
+                    : 'Not Loaded'}
              </Badge>
+              {health.model_downloaded === true && !health.model_loaded && (
+                <Badge variant="outline">
+                  Model Cached (will load on first use)
+                </Badge>
+              )}
              <Badge variant={health.gpu_available ? 'default' : 'secondary'}>
                GPU: {health.gpu_available ? 'Available' : 'Not Available'}
              </Badge>