mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-15 04:40:40 -07:00
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.
This commit is contained in:
Jamie Pine
+38
-50
@@ -1,63 +1,51 @@
|
||||
import { History, Mic, Settings, Sparkles } from 'lucide-react';
|
||||
import { useState } from 'react';
|
||||
import { GenerationForm } from '@/components/Generation/GenerationForm';
|
||||
import { HistoryTable } from '@/components/History/HistoryTable';
|
||||
import { ConnectionForm } from '@/components/ServerSettings/ConnectionForm';
|
||||
import { ServerStatus } from '@/components/ServerSettings/ServerStatus';
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
|
||||
import { ModelManagement } from '@/components/ServerSettings/ModelManagement';
|
||||
import { Toaster } from '@/components/ui/toaster';
|
||||
import { ProfileList } from '@/components/VoiceProfiles/ProfileList';
|
||||
import { Sidebar } from '@/components/Sidebar';
|
||||
|
||||
function App() {
|
||||
const [activeTab, setActiveTab] = useState('profiles');
|
||||
|
||||
return (
|
||||
<div className="min-h-screen bg-background">
|
||||
<div className="container mx-auto px-4 py-8 max-w-7xl">
|
||||
<div className="mb-8">
|
||||
<h1 className="text-4xl font-bold mb-2">voicebox</h1>
|
||||
<p className="text-muted-foreground">
|
||||
Production-quality Qwen3-TTS voice cloning and generation
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<Tabs defaultValue="profiles" className="space-y-4">
|
||||
<TabsList>
|
||||
<TabsTrigger value="profiles">
|
||||
<Mic className="mr-2 h-4 w-4" />
|
||||
Profiles
|
||||
</TabsTrigger>
|
||||
<TabsTrigger value="generate">
|
||||
<Sparkles className="mr-2 h-4 w-4" />
|
||||
Generate
|
||||
</TabsTrigger>
|
||||
<TabsTrigger value="history">
|
||||
<History className="mr-2 h-4 w-4" />
|
||||
History
|
||||
</TabsTrigger>
|
||||
<TabsTrigger value="settings">
|
||||
<Settings className="mr-2 h-4 w-4" />
|
||||
Settings
|
||||
</TabsTrigger>
|
||||
</TabsList>
|
||||
|
||||
<TabsContent value="profiles" className="space-y-4">
|
||||
<ProfileList />
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="generate" className="space-y-4">
|
||||
<GenerationForm />
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="history" className="space-y-4">
|
||||
<HistoryTable />
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="settings" className="space-y-4">
|
||||
<div className="grid gap-4 md:grid-cols-2">
|
||||
<ConnectionForm />
|
||||
<ServerStatus />
|
||||
<div className="min-h-screen bg-background flex">
|
||||
<Sidebar activeTab={activeTab} onTabChange={setActiveTab} />
|
||||
|
||||
<main className="flex-1 ml-20">
|
||||
<div className="container mx-auto px-8 py-8 max-w-7xl">
|
||||
{activeTab === 'profiles' && (
|
||||
<div className="space-y-4">
|
||||
<ProfileList />
|
||||
</div>
|
||||
</TabsContent>
|
||||
</Tabs>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{activeTab === 'generate' && (
|
||||
<div className="space-y-4">
|
||||
<GenerationForm />
|
||||
</div>
|
||||
)}
|
||||
|
||||
{activeTab === 'history' && (
|
||||
<div className="space-y-4">
|
||||
<HistoryTable />
|
||||
</div>
|
||||
)}
|
||||
|
||||
{activeTab === 'settings' && (
|
||||
<div className="space-y-4">
|
||||
<div className="grid gap-4 md:grid-cols-2">
|
||||
<ConnectionForm />
|
||||
<ServerStatus />
|
||||
</div>
|
||||
<ModelManagement />
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</main>
|
||||
|
||||
<Toaster />
|
||||
</div>
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
import { History, Mic, Settings, Sparkles } from 'lucide-react';
|
||||
import { cn } from '@/lib/utils/cn';
|
||||
|
||||
interface SidebarProps {
|
||||
activeTab: string;
|
||||
onTabChange: (tab: string) => void;
|
||||
}
|
||||
|
||||
const tabs = [
|
||||
{ id: 'profiles', icon: Mic, label: 'Profiles' },
|
||||
{ id: 'generate', icon: Sparkles, label: 'Generate' },
|
||||
{ id: 'history', icon: History, label: 'History' },
|
||||
{ id: 'settings', icon: Settings, label: 'Settings' },
|
||||
];
|
||||
|
||||
export function Sidebar({ activeTab, onTabChange }: SidebarProps) {
|
||||
return (
|
||||
<div className="fixed left-0 top-0 h-full w-20 bg-sidebar border-r border-border flex flex-col items-center py-6 gap-6">
|
||||
{/* Navigation Buttons */}
|
||||
<div className="flex flex-col gap-3">
|
||||
{tabs.map((tab) => {
|
||||
const Icon = tab.icon;
|
||||
const isActive = activeTab === tab.id;
|
||||
|
||||
return (
|
||||
<button
|
||||
key={tab.id}
|
||||
type="button"
|
||||
onClick={() => onTabChange(tab.id)}
|
||||
className={cn(
|
||||
"w-12 h-12 rounded-full flex items-center justify-center transition-all duration-200",
|
||||
"hover:bg-accent hover:text-accent-foreground",
|
||||
isActive
|
||||
? "bg-primary text-primary-foreground shadow-lg"
|
||||
: "text-muted-foreground"
|
||||
)}
|
||||
title={tab.label}
|
||||
aria-label={tab.label}
|
||||
>
|
||||
<Icon className="h-5 w-5" />
|
||||
</button>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
+30
-19
@@ -32,6 +32,7 @@
|
||||
--color-border: hsl(var(--border));
|
||||
--color-input: hsl(var(--input));
|
||||
--color-ring: hsl(var(--ring));
|
||||
--color-sidebar: hsl(var(--sidebar));
|
||||
|
||||
--color-chart-1: hsl(var(--chart-1));
|
||||
--color-chart-2: hsl(var(--chart-2));
|
||||
@@ -60,6 +61,7 @@
|
||||
--border: 214.3 31.8% 91.4%;
|
||||
--input: 214.3 31.8% 91.4%;
|
||||
--ring: 222.2 84% 4.9%;
|
||||
--sidebar: 0 0% 98%;
|
||||
--radius: 0.5rem;
|
||||
--chart-1: 12 76% 61%;
|
||||
--chart-2: 173 58% 39%;
|
||||
@@ -69,25 +71,26 @@
|
||||
}
|
||||
|
||||
.dark {
|
||||
--background: 222.2 84% 4.9%;
|
||||
--foreground: 210 40% 98%;
|
||||
--card: 222.2 84% 4.9%;
|
||||
--card-foreground: 210 40% 98%;
|
||||
--popover: 222.2 84% 4.9%;
|
||||
--popover-foreground: 210 40% 98%;
|
||||
--primary: 210 40% 98%;
|
||||
--primary-foreground: 222.2 47.4% 11.2%;
|
||||
--secondary: 217.2 32.6% 17.5%;
|
||||
--secondary-foreground: 210 40% 98%;
|
||||
--muted: 217.2 32.6% 17.5%;
|
||||
--muted-foreground: 215 20.2% 65.1%;
|
||||
--accent: 217.2 32.6% 17.5%;
|
||||
--accent-foreground: 210 40% 98%;
|
||||
--destructive: 0 62.8% 30.6%;
|
||||
--destructive-foreground: 210 40% 98%;
|
||||
--border: 217.2 32.6% 17.5%;
|
||||
--input: 217.2 32.6% 17.5%;
|
||||
--ring: 212.7 26.8% 83.9%;
|
||||
--background: 0 0% 8%;
|
||||
--foreground: 0 0% 95%;
|
||||
--card: 0 0% 10%;
|
||||
--card-foreground: 0 0% 95%;
|
||||
--popover: 0 0% 10%;
|
||||
--popover-foreground: 0 0% 95%;
|
||||
--primary: 0 0% 20%;
|
||||
--primary-foreground: 0 0% 95%;
|
||||
--secondary: 0 0% 15%;
|
||||
--secondary-foreground: 0 0% 95%;
|
||||
--muted: 0 0% 15%;
|
||||
--muted-foreground: 0 0% 60%;
|
||||
--accent: 0 0% 15%;
|
||||
--accent-foreground: 0 0% 95%;
|
||||
--destructive: 0 62.8% 50%;
|
||||
--destructive-foreground: 0 0% 95%;
|
||||
--border: 0 0% 15%;
|
||||
--input: 0 0% 15%;
|
||||
--ring: 0 0% 40%;
|
||||
--sidebar: 0 0% 6%;
|
||||
--chart-1: 220 70% 50%;
|
||||
--chart-2: 160 60% 45%;
|
||||
--chart-3: 30 80% 55%;
|
||||
@@ -103,3 +106,11 @@
|
||||
@apply bg-background text-foreground;
|
||||
}
|
||||
}
|
||||
|
||||
@layer utilities {
|
||||
.writing-vertical {
|
||||
writing-mode: vertical-rl;
|
||||
text-orientation: mixed;
|
||||
letter-spacing: 0.1em;
|
||||
}
|
||||
}
|
||||
|
||||
+290
-1
@@ -19,6 +19,7 @@ import uuid
|
||||
|
||||
from . import database, models, profiles, history, tts, transcribe
|
||||
from .database import get_db, init_db, Generation as DBGeneration, VoiceProfile as DBVoiceProfile
|
||||
from .utils.progress import get_progress_manager
|
||||
|
||||
# Initialize database
|
||||
init_db()
|
||||
@@ -52,6 +53,10 @@ async def root():
|
||||
@app.get("/health", response_model=models.HealthResponse)
|
||||
async def health():
|
||||
"""Health check endpoint."""
|
||||
from huggingface_hub import hf_hub_download
|
||||
from pathlib import Path
|
||||
import os
|
||||
|
||||
tts_model = tts.get_tts_model()
|
||||
gpu_available = torch.cuda.is_available()
|
||||
|
||||
@@ -59,9 +64,44 @@ async def health():
|
||||
if gpu_available:
|
||||
vram_used = torch.cuda.memory_allocated() / 1024 / 1024 # MB
|
||||
|
||||
# Check if model is loaded
|
||||
model_loaded = tts_model.is_loaded()
|
||||
model_size = tts_model.model_size if model_loaded else None
|
||||
|
||||
# Check if default model is downloaded (cached)
|
||||
model_downloaded = None
|
||||
try:
|
||||
# Check if the default model (1.7B) is cached
|
||||
default_model_id = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
|
||||
|
||||
# Method 1: Try scan_cache_dir if available
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
cache_info = scan_cache_dir()
|
||||
for repo in cache_info.repos:
|
||||
if repo.repo_id == default_model_id:
|
||||
model_downloaded = True
|
||||
break
|
||||
except (ImportError, Exception):
|
||||
# Method 2: Check cache directory
|
||||
cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
|
||||
repo_cache = Path(cache_dir) / "models--" + default_model_id.replace("/", "--")
|
||||
if repo_cache.exists():
|
||||
has_model_files = (
|
||||
any(repo_cache.rglob("*.bin")) or
|
||||
any(repo_cache.rglob("*.safetensors")) or
|
||||
any(repo_cache.rglob("*.pt")) or
|
||||
any(repo_cache.rglob("*.pth"))
|
||||
)
|
||||
model_downloaded = has_model_files
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return models.HealthResponse(
|
||||
status="healthy",
|
||||
model_loaded=tts_model.is_loaded(),
|
||||
model_loaded=model_loaded,
|
||||
model_downloaded=model_downloaded,
|
||||
model_size=model_size,
|
||||
gpu_available=gpu_available,
|
||||
vram_used_mb=vram_used,
|
||||
)
|
||||
@@ -395,6 +435,255 @@ async def unload_model():
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.get("/models/progress/{model_name}")
|
||||
async def get_model_progress(model_name: str):
|
||||
"""Get model download progress via Server-Sent Events."""
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
progress_manager = get_progress_manager()
|
||||
|
||||
async def event_generator():
|
||||
"""Generate SSE events for progress updates."""
|
||||
async for event in progress_manager.subscribe(model_name):
|
||||
yield event
|
||||
|
||||
return StreamingResponse(
|
||||
event_generator(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@app.get("/models/status", response_model=models.ModelStatusListResponse)
|
||||
async def get_model_status():
|
||||
"""Get status of all available models."""
|
||||
from huggingface_hub import hf_hub_download
|
||||
from pathlib import Path
|
||||
import os
|
||||
|
||||
# Try to import scan_cache_dir (might not be available in older versions)
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
use_scan_cache = True
|
||||
except ImportError:
|
||||
use_scan_cache = False
|
||||
|
||||
def check_tts_loaded(model_size: str):
|
||||
"""Check if TTS model is loaded with specific size."""
|
||||
try:
|
||||
tts_model = tts.get_tts_model()
|
||||
return tts_model.is_loaded() and tts_model.model_size == model_size
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def check_whisper_loaded(model_size: str):
|
||||
"""Check if Whisper model is loaded with specific size."""
|
||||
try:
|
||||
whisper_model = transcribe.get_whisper_model()
|
||||
return whisper_model.is_loaded() and whisper_model.model_size == model_size
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
model_configs = [
|
||||
{
|
||||
"model_name": "qwen-tts-1.7B",
|
||||
"display_name": "Qwen TTS 1.7B",
|
||||
"hf_repo_id": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
|
||||
"model_size": "1.7B",
|
||||
"check_loaded": lambda: check_tts_loaded("1.7B"),
|
||||
},
|
||||
{
|
||||
"model_name": "qwen-tts-0.6B",
|
||||
"display_name": "Qwen TTS 0.6B",
|
||||
"hf_repo_id": "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
|
||||
"model_size": "0.6B",
|
||||
"check_loaded": lambda: check_tts_loaded("0.6B"),
|
||||
},
|
||||
{
|
||||
"model_name": "whisper-base",
|
||||
"display_name": "Whisper Base",
|
||||
"hf_repo_id": "openai/whisper-base",
|
||||
"model_size": "base",
|
||||
"check_loaded": lambda: check_whisper_loaded("base"),
|
||||
},
|
||||
{
|
||||
"model_name": "whisper-small",
|
||||
"display_name": "Whisper Small",
|
||||
"hf_repo_id": "openai/whisper-small",
|
||||
"model_size": "small",
|
||||
"check_loaded": lambda: check_whisper_loaded("small"),
|
||||
},
|
||||
{
|
||||
"model_name": "whisper-medium",
|
||||
"display_name": "Whisper Medium",
|
||||
"hf_repo_id": "openai/whisper-medium",
|
||||
"model_size": "medium",
|
||||
"check_loaded": lambda: check_whisper_loaded("medium"),
|
||||
},
|
||||
{
|
||||
"model_name": "whisper-large",
|
||||
"display_name": "Whisper Large",
|
||||
"hf_repo_id": "openai/whisper-large",
|
||||
"model_size": "large",
|
||||
"check_loaded": lambda: check_whisper_loaded("large"),
|
||||
},
|
||||
]
|
||||
|
||||
# Get HuggingFace cache info (if available)
|
||||
cache_info = None
|
||||
if use_scan_cache:
|
||||
try:
|
||||
cache_info = scan_cache_dir()
|
||||
except Exception:
|
||||
# Function failed, continue without it
|
||||
pass
|
||||
|
||||
statuses = []
|
||||
|
||||
for config in model_configs:
|
||||
try:
|
||||
downloaded = False
|
||||
size_mb = None
|
||||
loaded = False
|
||||
|
||||
# Method 1: Try using scan_cache_dir if available
|
||||
if cache_info:
|
||||
repo_id = config["hf_repo_id"]
|
||||
for repo in cache_info.repos:
|
||||
if repo.repo_id == repo_id:
|
||||
downloaded = True
|
||||
# Calculate size from cache info
|
||||
try:
|
||||
total_size = sum(revision.size_on_disk for revision in repo.revisions)
|
||||
size_mb = total_size / (1024 * 1024)
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
|
||||
# Method 2: Fallback to checking cache directory directly
|
||||
if not downloaded:
|
||||
try:
|
||||
cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
|
||||
repo_cache = Path(cache_dir) / "models--" + config["hf_repo_id"].replace("/", "--")
|
||||
|
||||
if repo_cache.exists():
|
||||
# Check for model files (bin, safetensors, or other common model files)
|
||||
has_model_files = (
|
||||
any(repo_cache.rglob("*.bin")) or
|
||||
any(repo_cache.rglob("*.safetensors")) or
|
||||
any(repo_cache.rglob("*.pt")) or
|
||||
any(repo_cache.rglob("*.pth")) or
|
||||
any(repo_cache.rglob("model.safetensors.index.json")) or
|
||||
any(repo_cache.rglob("pytorch_model.bin.index.json"))
|
||||
)
|
||||
|
||||
if has_model_files:
|
||||
downloaded = True
|
||||
# Calculate size
|
||||
try:
|
||||
total_size = sum(f.stat().st_size for f in repo_cache.rglob("*") if f.is_file())
|
||||
size_mb = total_size / (1024 * 1024)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Method 3: Try to check if model can be loaded locally (last resort)
|
||||
if not downloaded:
|
||||
try:
|
||||
# Try to download with local_files_only=True to check if cached
|
||||
hf_hub_download(
|
||||
repo_id=config["hf_repo_id"],
|
||||
filename="config.json", # Try a common file
|
||||
local_files_only=True,
|
||||
)
|
||||
downloaded = True
|
||||
except Exception:
|
||||
# File not found locally, model not downloaded
|
||||
pass
|
||||
|
||||
# Check if loaded in memory
|
||||
try:
|
||||
loaded = config["check_loaded"]()
|
||||
except Exception:
|
||||
loaded = False
|
||||
|
||||
statuses.append(models.ModelStatus(
|
||||
model_name=config["model_name"],
|
||||
display_name=config["display_name"],
|
||||
downloaded=downloaded,
|
||||
size_mb=size_mb,
|
||||
loaded=loaded,
|
||||
))
|
||||
except Exception as e:
|
||||
# If check fails, try to at least check if loaded
|
||||
try:
|
||||
loaded = config["check_loaded"]()
|
||||
except Exception:
|
||||
loaded = False
|
||||
|
||||
statuses.append(models.ModelStatus(
|
||||
model_name=config["model_name"],
|
||||
display_name=config["display_name"],
|
||||
downloaded=False, # Assume not downloaded if check failed
|
||||
size_mb=None,
|
||||
loaded=loaded,
|
||||
))
|
||||
|
||||
return models.ModelStatusListResponse(models=statuses)
|
||||
|
||||
|
||||
@app.post("/models/download")
|
||||
async def trigger_model_download(request: models.ModelDownloadRequest):
|
||||
"""Trigger download of a specific model."""
|
||||
import asyncio
|
||||
|
||||
model_configs = {
|
||||
"qwen-tts-1.7B": {
|
||||
"model_size": "1.7B",
|
||||
"load_func": lambda: tts.get_tts_model().load_model("1.7B"),
|
||||
},
|
||||
"qwen-tts-0.6B": {
|
||||
"model_size": "0.6B",
|
||||
"load_func": lambda: tts.get_tts_model().load_model("0.6B"),
|
||||
},
|
||||
"whisper-base": {
|
||||
"model_size": "base",
|
||||
"load_func": lambda: transcribe.get_whisper_model().load_model("base"),
|
||||
},
|
||||
"whisper-small": {
|
||||
"model_size": "small",
|
||||
"load_func": lambda: transcribe.get_whisper_model().load_model("small"),
|
||||
},
|
||||
"whisper-medium": {
|
||||
"model_size": "medium",
|
||||
"load_func": lambda: transcribe.get_whisper_model().load_model("medium"),
|
||||
},
|
||||
"whisper-large": {
|
||||
"model_size": "large",
|
||||
"load_func": lambda: transcribe.get_whisper_model().load_model("large"),
|
||||
},
|
||||
}
|
||||
|
||||
if request.model_name not in model_configs:
|
||||
raise HTTPException(status_code=400, detail=f"Unknown model: {request.model_name}")
|
||||
|
||||
config = model_configs[request.model_name]
|
||||
|
||||
try:
|
||||
# Trigger download by loading the model (which will download if not cached)
|
||||
# Run in background to avoid blocking
|
||||
await asyncio.to_thread(config["load_func"])
|
||||
|
||||
return {"message": f"Model {request.model_name} download started"}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ============================================
|
||||
# STARTUP & SHUTDOWN
|
||||
# ============================================
|
||||
|
||||
@@ -111,5 +111,26 @@ class HealthResponse(BaseModel):
|
||||
"""Response model for health check."""
|
||||
status: str
|
||||
model_loaded: bool
|
||||
model_downloaded: Optional[bool] = None # Whether model is cached/downloaded
|
||||
model_size: Optional[str] = None # Current model size if loaded
|
||||
gpu_available: bool
|
||||
vram_used_mb: Optional[float] = None
|
||||
|
||||
|
||||
class ModelStatus(BaseModel):
|
||||
"""Response model for model status."""
|
||||
model_name: str
|
||||
display_name: str
|
||||
downloaded: bool
|
||||
size_mb: Optional[float] = None
|
||||
loaded: bool = False
|
||||
|
||||
|
||||
class ModelStatusListResponse(BaseModel):
|
||||
"""Response model for model status list."""
|
||||
models: List[ModelStatus]
|
||||
|
||||
|
||||
class ModelDownloadRequest(BaseModel):
|
||||
"""Request model for triggering model download."""
|
||||
model_name: str
|
||||
|
||||
+20
-3
@@ -6,6 +6,8 @@ from typing import Optional, List, Dict
|
||||
import torch
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from .utils.progress import get_progress_manager
|
||||
from .utils.hf_progress import HFProgressTracker, create_hf_progress_callback
|
||||
|
||||
|
||||
class WhisperModel:
|
||||
@@ -48,18 +50,33 @@ class WhisperModel:
|
||||
|
||||
model_name = f"openai/whisper-{model_size}"
|
||||
|
||||
# Set up progress tracking
|
||||
progress_manager = get_progress_manager()
|
||||
progress_model_name = f"whisper-{model_size}"
|
||||
|
||||
print(f"Loading Whisper model {model_size} on {self.device}...")
|
||||
|
||||
self.processor = WhisperProcessor.from_pretrained(model_name)
|
||||
self.model = WhisperForConditionalGeneration.from_pretrained(model_name)
|
||||
self.model.to(self.device)
|
||||
# Set up progress callback
|
||||
progress_callback = create_hf_progress_callback(progress_model_name, progress_manager)
|
||||
tracker = HFProgressTracker(progress_callback)
|
||||
|
||||
# Use progress tracker during download
|
||||
with tracker.patch_download():
|
||||
self.processor = WhisperProcessor.from_pretrained(model_name)
|
||||
self.model = WhisperForConditionalGeneration.from_pretrained(model_name)
|
||||
|
||||
self.model.to(self.device)
|
||||
self.model_size = model_size
|
||||
|
||||
# Mark as complete
|
||||
progress_manager.mark_complete(progress_model_name)
|
||||
|
||||
print(f"Whisper model {model_size} loaded successfully")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error loading Whisper model: {e}")
|
||||
progress_manager = get_progress_manager()
|
||||
progress_manager.mark_error(f"whisper-{model_size}", str(e))
|
||||
raise
|
||||
|
||||
def unload_model(self):
|
||||
|
||||
+36
-7
@@ -11,6 +11,8 @@ from pathlib import Path
|
||||
|
||||
from .utils.cache import get_cache_key, get_cached_voice_prompt, cache_voice_prompt
|
||||
from .utils.audio import normalize_audio
|
||||
from .utils.progress import get_progress_manager
|
||||
from .utils.hf_progress import HFProgressTracker, create_hf_progress_callback
|
||||
|
||||
|
||||
class TTSModel:
|
||||
@@ -99,14 +101,37 @@ class TTSModel:
|
||||
# Get model path (local or HuggingFace Hub ID)
|
||||
model_path = self._get_model_path(model_size)
|
||||
|
||||
print(f"Loading TTS model {model_size} on {self.device}...")
|
||||
# Set up progress tracking
|
||||
progress_manager = get_progress_manager()
|
||||
model_name = f"qwen-tts-{model_size}"
|
||||
|
||||
# Load the model - from_pretrained handles both local paths and HF Hub IDs
|
||||
self.model = Qwen3TTSModel.from_pretrained(
|
||||
model_path,
|
||||
device_map=self.device,
|
||||
torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16,
|
||||
)
|
||||
# Check if model is being downloaded from HuggingFace Hub
|
||||
if model_path.startswith("Qwen/"):
|
||||
print(f"Loading TTS model {model_size} on {self.device}...")
|
||||
|
||||
# Set up progress callback
|
||||
progress_callback = create_hf_progress_callback(model_name, progress_manager)
|
||||
tracker = HFProgressTracker(progress_callback)
|
||||
|
||||
# Use progress tracker during download
|
||||
with tracker.patch_download():
|
||||
# Load the model - downloads will happen automatically with progress tracking
|
||||
self.model = Qwen3TTSModel.from_pretrained(
|
||||
model_path,
|
||||
device_map=self.device,
|
||||
torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16,
|
||||
)
|
||||
|
||||
# Mark as complete
|
||||
progress_manager.mark_complete(model_name)
|
||||
else:
|
||||
# Local model, no download needed
|
||||
print(f"Loading TTS model {model_size} on {self.device}...")
|
||||
self.model = Qwen3TTSModel.from_pretrained(
|
||||
model_path,
|
||||
device_map=self.device,
|
||||
torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16,
|
||||
)
|
||||
|
||||
self._current_model_size = model_size
|
||||
self.model_size = model_size
|
||||
@@ -115,10 +140,14 @@ class TTSModel:
|
||||
|
||||
except ImportError as e:
|
||||
print(f"Error: qwen_tts package not found. Install with: pip install git+https://github.com/QwenLM/Qwen3-TTS.git")
|
||||
progress_manager = get_progress_manager()
|
||||
progress_manager.mark_error(f"qwen-tts-{model_size}", str(e))
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"Error loading TTS model: {e}")
|
||||
print(f"Tip: The model will be automatically downloaded from HuggingFace Hub on first use.")
|
||||
progress_manager = get_progress_manager()
|
||||
progress_manager.mark_error(f"qwen-tts-{model_size}", str(e))
|
||||
raise
|
||||
|
||||
def unload_model(self):
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
HuggingFace Hub download progress tracking.
|
||||
"""
|
||||
|
||||
from typing import Optional, Callable
|
||||
from contextlib import contextmanager
|
||||
import threading
|
||||
|
||||
|
||||
class HFProgressTracker:
|
||||
"""Tracks HuggingFace Hub download progress by intercepting hf_hub_download."""
|
||||
|
||||
def __init__(self, progress_callback: Optional[Callable] = None):
|
||||
self.progress_callback = progress_callback
|
||||
self._original_hf_hub_download = None
|
||||
self._lock = threading.Lock()
|
||||
self._total_downloaded = 0
|
||||
self._total_size = 0
|
||||
|
||||
def _tracked_hf_hub_download(self, *args, **kwargs):
|
||||
"""Wrapper for hf_hub_download with progress tracking."""
|
||||
import huggingface_hub
|
||||
|
||||
# Get original callback if present
|
||||
original_resume_callback = kwargs.get("resume_download", None)
|
||||
|
||||
def combined_callback(downloaded: int, total: int):
|
||||
"""Combined callback that tracks progress."""
|
||||
# Update totals
|
||||
with self._lock:
|
||||
# Estimate: assume each file contributes equally
|
||||
# This is a simplification - in reality we'd track per-file
|
||||
if total > 0:
|
||||
self._total_size = max(self._total_size, total)
|
||||
self._total_downloaded = downloaded
|
||||
|
||||
# Call original callback if present
|
||||
if original_resume_callback:
|
||||
original_resume_callback(downloaded, total)
|
||||
|
||||
# Call our progress callback
|
||||
if self.progress_callback:
|
||||
with self._lock:
|
||||
self.progress_callback(self._total_downloaded, self._total_size)
|
||||
|
||||
# Replace callback
|
||||
kwargs["resume_download"] = combined_callback
|
||||
|
||||
# Call original download
|
||||
return self._original_hf_hub_download(*args, **kwargs)
|
||||
|
||||
@contextmanager
|
||||
def patch_download(self):
|
||||
"""Context manager to patch hf_hub_download for progress tracking."""
|
||||
try:
|
||||
import huggingface_hub
|
||||
self._original_hf_hub_download = huggingface_hub.hf_hub_download
|
||||
|
||||
# Reset totals
|
||||
with self._lock:
|
||||
self._total_downloaded = 0
|
||||
self._total_size = 0
|
||||
|
||||
# Patch the function
|
||||
huggingface_hub.hf_hub_download = self._tracked_hf_hub_download
|
||||
|
||||
yield
|
||||
except ImportError:
|
||||
# If huggingface_hub not available, just yield without patching
|
||||
yield
|
||||
finally:
|
||||
# Restore original
|
||||
if self._original_hf_hub_download:
|
||||
try:
|
||||
import huggingface_hub
|
||||
huggingface_hub.hf_hub_download = self._original_hf_hub_download
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def create_hf_progress_callback(model_name: str, progress_manager):
|
||||
"""Create a progress callback for HuggingFace downloads."""
|
||||
def callback(downloaded: int, total: int):
|
||||
"""Progress callback."""
|
||||
if total > 0:
|
||||
progress_manager.update_progress(
|
||||
model_name=model_name,
|
||||
current=downloaded,
|
||||
total=total,
|
||||
filename="",
|
||||
status="downloading",
|
||||
)
|
||||
return callback
|
||||
@@ -0,0 +1,164 @@
|
||||
"""
|
||||
Progress tracking for model downloads using Server-Sent Events.
|
||||
"""
|
||||
|
||||
from typing import Optional, Callable, Dict
|
||||
from fastapi.responses import StreamingResponse
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class ProgressManager:
|
||||
"""Manages download progress for multiple models."""
|
||||
|
||||
def __init__(self):
|
||||
self._progress: Dict[str, Dict] = {}
|
||||
self._listeners: Dict[str, list] = {}
|
||||
|
||||
def update_progress(
|
||||
self,
|
||||
model_name: str,
|
||||
current: int,
|
||||
total: int,
|
||||
filename: Optional[str] = None,
|
||||
status: str = "downloading",
|
||||
):
|
||||
"""
|
||||
Update progress for a model download.
|
||||
|
||||
Args:
|
||||
model_name: Name of the model (e.g., "qwen-tts-1.7B", "whisper-base")
|
||||
current: Current bytes downloaded
|
||||
total: Total bytes to download
|
||||
filename: Current file being downloaded
|
||||
status: Status string (downloading, extracting, complete, error)
|
||||
"""
|
||||
progress_pct = (current / total * 100) if total > 0 else 0
|
||||
|
||||
self._progress[model_name] = {
|
||||
"model_name": model_name,
|
||||
"current": current,
|
||||
"total": total,
|
||||
"progress": progress_pct,
|
||||
"filename": filename,
|
||||
"status": status,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
}
|
||||
|
||||
# Notify all listeners
|
||||
if model_name in self._listeners:
|
||||
for queue in self._listeners[model_name]:
|
||||
try:
|
||||
queue.put_nowait(self._progress[model_name].copy())
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
|
||||
def get_progress(self, model_name: str) -> Optional[Dict]:
|
||||
"""Get current progress for a model."""
|
||||
return self._progress.get(model_name)
|
||||
|
||||
def create_progress_callback(self, model_name: str, filename: Optional[str] = None):
|
||||
"""
|
||||
Create a progress callback function for HuggingFace downloads.
|
||||
|
||||
Args:
|
||||
model_name: Name of the model
|
||||
filename: Optional filename filter
|
||||
|
||||
Returns:
|
||||
Callback function
|
||||
"""
|
||||
def callback(progress: Dict):
|
||||
"""HuggingFace Hub progress callback."""
|
||||
if "total" in progress and "current" in progress:
|
||||
current = progress.get("current", 0)
|
||||
total = progress.get("total", 0)
|
||||
file_name = progress.get("filename", filename)
|
||||
|
||||
self.update_progress(
|
||||
model_name=model_name,
|
||||
current=current,
|
||||
total=total,
|
||||
filename=file_name,
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
return callback
|
||||
|
||||
async def subscribe(self, model_name: str):
|
||||
"""
|
||||
Subscribe to progress updates for a model.
|
||||
|
||||
Yields progress updates as Server-Sent Events.
|
||||
"""
|
||||
queue = asyncio.Queue(maxsize=10)
|
||||
|
||||
# Add to listeners
|
||||
if model_name not in self._listeners:
|
||||
self._listeners[model_name] = []
|
||||
self._listeners[model_name].append(queue)
|
||||
|
||||
try:
|
||||
# Send initial progress if available
|
||||
if model_name in self._progress:
|
||||
yield f"data: {json.dumps(self._progress[model_name])}\n\n"
|
||||
|
||||
# Stream updates
|
||||
while True:
|
||||
try:
|
||||
# Wait for update with timeout
|
||||
progress = await asyncio.wait_for(queue.get(), timeout=1.0)
|
||||
yield f"data: {json.dumps(progress)}\n\n"
|
||||
|
||||
# Stop if complete or error
|
||||
if progress.get("status") in ("complete", "error"):
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
# Send heartbeat
|
||||
yield ": heartbeat\n\n"
|
||||
continue
|
||||
finally:
|
||||
# Remove from listeners
|
||||
if model_name in self._listeners:
|
||||
self._listeners[model_name].remove(queue)
|
||||
if not self._listeners[model_name]:
|
||||
del self._listeners[model_name]
|
||||
|
||||
def mark_complete(self, model_name: str):
|
||||
"""Mark a model download as complete."""
|
||||
if model_name in self._progress:
|
||||
self._progress[model_name]["status"] = "complete"
|
||||
self._progress[model_name]["progress"] = 100.0
|
||||
# Notify listeners
|
||||
if model_name in self._listeners:
|
||||
for queue in self._listeners[model_name]:
|
||||
try:
|
||||
queue.put_nowait(self._progress[model_name].copy())
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
|
||||
def mark_error(self, model_name: str, error: str):
|
||||
"""Mark a model download as failed."""
|
||||
if model_name in self._progress:
|
||||
self._progress[model_name]["status"] = "error"
|
||||
self._progress[model_name]["error"] = error
|
||||
# Notify listeners
|
||||
if model_name in self._listeners:
|
||||
for queue in self._listeners[model_name]:
|
||||
try:
|
||||
queue.put_nowait(self._progress[model_name].copy())
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
|
||||
|
||||
# Global progress manager instance
|
||||
_progress_manager: Optional[ProgressManager] = None
|
||||
|
||||
|
||||
def get_progress_manager() -> ProgressManager:
|
||||
"""Get or create the global progress manager."""
|
||||
global _progress_manager
|
||||
if _progress_manager is None:
|
||||
_progress_manager = ProgressManager()
|
||||
return _progress_manager
|
||||
Reference in New Issue
Block a user