| |
|
| | import { useState, useCallback, useRef } from 'react'; |
| | import { FileItem, UploadStatus } from '../types'; |
| | import { uploadBatchToHub } from '../services/hfService'; |
| |
|
| | export const useFileUpload = () => { |
| | const [files, setFiles] = useState<FileItem[]>([]); |
| | const [isUploading, setIsUploading] = useState(false); |
| | |
| | |
| | const filesRef = useRef<FileItem[]>([]); |
| | filesRef.current = files; |
| |
|
| | |
| | |
| | const BATCH_SIZE = 10; |
| | |
| | const CONCURRENCY_LIMIT = 5; |
| |
|
| | |
| |
|
| | const addFiles = useCallback((newFilesList: FileItem[]) => { |
| | setFiles((prev) => [...prev, ...newFilesList]); |
| | }, []); |
| |
|
| | const removeFile = useCallback((id: string) => { |
| | setFiles((prev) => prev.filter((f) => f.id !== id)); |
| | }, []); |
| |
|
| | const updateFilePath = useCallback((id: string, newPath: string) => { |
| | setFiles((prev) => prev.map((f) => (f.id === id ? { ...f, path: newPath } : f))); |
| | }, []); |
| |
|
| | const startUpload = useCallback(async () => { |
| | |
| | const pendingFiles = filesRef.current.filter( |
| | (f) => f.status === UploadStatus.IDLE || f.status === UploadStatus.ERROR |
| | ); |
| |
|
| | if (pendingFiles.length === 0) return; |
| |
|
| | setIsUploading(true); |
| |
|
| | |
| | const batches: FileItem[][] = []; |
| | for (let i = 0; i < pendingFiles.length; i += BATCH_SIZE) { |
| | batches.push(pendingFiles.slice(i, i + BATCH_SIZE)); |
| | } |
| |
|
| | |
| | const queue = [...batches]; |
| | let activeWorkers = 0; |
| |
|
| | |
| | const updateBatchStatus = (batchItems: FileItem[], status: UploadStatus, result?: { urls?: string[], error?: string }) => { |
| | setFiles((prev) => |
| | prev.map((f) => { |
| | const batchIndex = batchItems.findIndex(b => b.id === f.id); |
| | if (batchIndex !== -1) { |
| | return { |
| | ...f, |
| | status: status, |
| | url: status === UploadStatus.SUCCESS ? result?.urls?.[batchIndex] : f.url, |
| | error: status === UploadStatus.ERROR ? result?.error : undefined |
| | }; |
| | } |
| | return f; |
| | }) |
| | ); |
| | }; |
| |
|
| | |
| | |
| | const processNextBatch = async (): Promise<void> => { |
| | if (queue.length === 0) return; |
| |
|
| | const batch = queue.shift(); |
| | if (!batch) return; |
| |
|
| | activeWorkers++; |
| | |
| | |
| | updateBatchStatus(batch, UploadStatus.UPLOADING); |
| |
|
| | try { |
| | const payload = batch.map(item => ({ |
| | id: item.id, |
| | file: item.file, |
| | path: item.path |
| | })); |
| |
|
| | const urls = await uploadBatchToHub(payload); |
| | |
| | |
| | updateBatchStatus(batch, UploadStatus.SUCCESS, { urls }); |
| |
|
| | } catch (err: any) { |
| | console.error("Batch failed:", err); |
| | |
| | updateBatchStatus(batch, UploadStatus.ERROR, { error: err.message || "Upload failed" }); |
| | } finally { |
| | activeWorkers--; |
| | |
| | if (queue.length > 0) { |
| | await processNextBatch(); |
| | } |
| | } |
| | }; |
| |
|
| | |
| | |
| | const initialWorkers = []; |
| | const limit = Math.min(CONCURRENCY_LIMIT, batches.length); |
| | |
| | for (let i = 0; i < limit; i++) { |
| | initialWorkers.push(processNextBatch()); |
| | } |
| |
|
| | |
| | await Promise.all(initialWorkers); |
| | |
| | setIsUploading(false); |
| |
|
| | }, []); |
| |
|
| | return { |
| | files, |
| | isUploading, |
| | addFiles, |
| | removeFile, |
| | updateFilePath, |
| | startUpload |
| | }; |
| | }; |
| |
|