File size: 8,550 Bytes
0573fbf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72f7156
0573fbf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72f7156
0573fbf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72f7156
0573fbf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
import React, { useState, useEffect } from 'react';
import {
    Dialog,
    DialogTitle,
    DialogContent,
    DialogActions,
    Button,
    Typography,
    TextField,
    Box,
    CircularProgress,
    Stepper,
    Step,
    StepLabel,
    Link,
    Alert,
    LinearProgress
} from '@mui/material';
import api from '../api';
import { hfAuthDialogStyles } from '../theme';

const HfAuthDialog = ({ open, onClose, onModelsDownloaded }) => {
    const [activeStep, setActiveStep] = useState(0);
    const [missingModels, setMissingModels] = useState([]);
    const [checkingStatus, setCheckingStatus] = useState(true);
    const [token, setToken] = useState('');
    const [error, setError] = useState(null);
    const [isProcessing, setIsProcessing] = useState(false);
    const [downloadingModel, setDownloadingModel] = useState(null);

    const steps = ['Check required models', 'Authenticate', 'Download models'];

    useEffect(() => {
        if (open) {
            checkModelStatus();
        } else {
            setActiveStep(0);
            setError(null);
            setToken('');
            setMissingModels([]);
        }
    }, [open]);

    const checkModelStatus = async () => {
        setCheckingStatus(true);
        setError(null);
        try {
            const response = await api.get('/api/base-models/status');
            const models = response.data.base_models;
            const missing = Object.entries(models)
                .filter(([_, info]) => !info.downloaded)
                .map(([id, info]) => ({ id, ...info }));
            
            setMissingModels(missing);
            
            if (missing.length === 0) {
                setActiveStep(3);
            } else {
                setActiveStep(1);
            }
        } catch (err) {
            setError(err.response?.data?.error || err.message || 'Failed to check model status.');
        } finally {
            setCheckingStatus(false);
        }
    };

    const handleLogin = async () => {
        if (!token.trim()) {
            setError('Please enter a Hugging Face token.');
            return;
        }

        setIsProcessing(true);
        setError(null);
        try {
            await api.post('/api/hf-login', { token: token.trim() });

            setActiveStep(2);
            startDownloads();
        } catch (err) {
            setError(err.response?.data?.error || err.message || 'Authentication failed. Please check your token.');
            setIsProcessing(false);
        }
    };

    const startDownloads = async () => {
        try {
            for (const model of missingModels) {
                setDownloadingModel(model.name);
                await api.post(`/api/models/${model.id}/accept-terms`);
                await api.post(`/api/models/${model.id}/download`);
            }

            setActiveStep(3);
            if (onModelsDownloaded) {
                onModelsDownloaded();
            }
        } catch (err) {
            setError(err.response?.data?.error || err.message || 'Failed to download models.');
        } finally {
            setIsProcessing(false);
            setDownloadingModel(null);
        }
    };

    const handleClose = () => {
        if (isProcessing && activeStep === 2) {
            return;
        }
        onClose(activeStep === 3);
    };

    const getStepContent = (stepIndex) => {
        if (checkingStatus) {
            return (
                <Box sx={hfAuthDialogStyles.checkingBox}>
                    <CircularProgress sx={hfAuthDialogStyles.checkingProgress} />
                    <Typography>Checking model availability...</Typography>
                </Box>
            );
        }

        switch (stepIndex) {
            case 1:
                return (
                    <Box sx={hfAuthDialogStyles.authStepBox}>
                        <Typography variant="body1" paragraph>
                            Some required base models are missing. You need a Hugging Face access token to download them.
                        </Typography>
                        
                        <Typography variant="body2" color="textSecondary" paragraph>
                            Before proceeding, please ensure you have visited huggingface.co and accepted the terms of use for the required models (e.g., stabilityai/stable-audio-open-1.0).
                        </Typography>
                        
                        <TextField
                            fullWidth
                            label="Hugging Face Access Token"
                            type="password"
                            value={token}
                            onChange={(e) => setToken(e.target.value)}
                            placeholder="hf_xxxxxxxxxxxxxxxxxxx..."
                            margin="normal"
                            variant="outlined"
                        />
                        <Typography variant="caption" color="textSecondary">
                            You can get a token from{' '}
                            <Link href="https://huggingface.co/settings/tokens" target="_blank" rel="noopener">
                                your Hugging Face settings
                            </Link>. "Read" access to public gated repos is needed.
                        </Typography>
                    </Box>
                );
            case 2:
                return (
                    <Box sx={hfAuthDialogStyles.downloadStepBox}>
                        <Typography variant="h6" paragraph>
                            Downloading {downloadingModel}...
                        </Typography>
                        <LinearProgress sx={hfAuthDialogStyles.downloadProgress} />
                        <Typography variant="body2" color="textSecondary">
                            This may take several minutes depending on your connection speed. Do not close the application.
                        </Typography>
                    </Box>
                );
            case 3:
                return (
                    <Box sx={hfAuthDialogStyles.successStepBox}>
                        <Typography variant="h6" color="success.main" paragraph>
                            All models are ready!
                        </Typography>
                        <Typography variant="body1">
                            You can now close this dialog and begin using Fragmenta.
                        </Typography>
                    </Box>
                );
            default:
                return "Unknown step";
        }
    };

    return (
        <Dialog 
            open={open} 
            onClose={handleClose}
            maxWidth="sm"
            fullWidth
            disableEscapeKeyDown={isProcessing && activeStep === 2}
        >
            <DialogTitle>
                Hugging Face Authentication
            </DialogTitle>
            <DialogContent dividers>
                <Stepper activeStep={activeStep} alternativeLabel sx={hfAuthDialogStyles.stepper}>
                    {steps.map((label) => (
                        <Step key={label}>
                            <StepLabel>{label}</StepLabel>
                        </Step>
                    ))}
                </Stepper>
                
                {error && (
                    <Alert severity="error" sx={hfAuthDialogStyles.errorAlert}>{error}</Alert>
                )}
                
                {getStepContent(activeStep)}
            </DialogContent>
            
            <DialogActions>
                {activeStep !== 3 && activeStep !== 2 && (
                    <Button onClick={handleClose} disabled={isProcessing}>
                        Cancel
                    </Button>
                )}
                
                {activeStep === 1 && (
                    <Button 
                        variant="contained" 
                        color="primary" 
                        onClick={handleLogin}
                        disabled={isProcessing || !token}
                    >
                        {isProcessing ? <CircularProgress size={hfAuthDialogStyles.loginSpinnerSize} /> : 'Login & Download'}
                    </Button>
                )}
                
                {activeStep === 3 && (
                    <Button variant="contained" color="primary" onClick={handleClose}>
                        Close
                    </Button>
                )}
            </DialogActions>
        </Dialog>
    );
};

export default HfAuthDialog;