Spaces:
Running on Zero
Running on Zero
File size: 2,157 Bytes
b701455 | 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 | import { useEffect } from 'react';
import { getSettingsPreferences, listControlNets, listModels } from '../api/client';
import { useStore } from '../store/useStore';
import { getDefaultModel, getModelSelectionUpdates } from '../lib/settings';
import { useShallow } from 'zustand/react/shallow';
export function useGenerationBootstrap() {
const { availableModels, setModels, setControlNets, setSettings } = useStore(useShallow((state) => ({
availableModels: state.availableModels,
setModels: state.setModels,
setControlNets: state.setControlNets,
setSettings: state.setSettings,
})));
useEffect(() => {
let cancelled = false;
const loadModels = async () => {
try {
const models = await listModels();
if (cancelled) return;
setModels(models);
const currentSettings = useStore.getState().settings;
if (!currentSettings.model_path && models.length > 0) {
const defaultModel = getDefaultModel(models);
if (defaultModel) {
setSettings(getModelSelectionUpdates(defaultModel, models));
}
}
} catch (error) {
console.error('Failed to load models', error);
}
};
const loadControlNets = async () => {
try {
const response = await listControlNets();
if (!cancelled) {
setControlNets(response.models);
}
} catch (error) {
console.error('Failed to load ControlNet models', error);
}
};
const loadPreferences = async () => {
try {
const preferences = await getSettingsPreferences();
if (!cancelled) {
setSettings({
torch_compile: !!preferences.torch_compile,
vae_autotune: !!preferences.vae_autotune,
});
}
} catch (error) {
console.error('Failed to load settings preferences', error);
}
};
if (availableModels.length === 0) {
void loadModels();
}
void loadControlNets();
void loadPreferences();
return () => {
cancelled = true;
};
}, [availableModels.length, setControlNets, setModels, setSettings]);
}
|