import { IconButton } from "@/components/ui/button" import { useToggle } from "@uidotdev/usehooks" import { Dialog, DialogContent, DialogTitle, DialogTrigger } from "./ui/dialog" import { HelpCircle, Settings } from "lucide-react" import { zodResolver } from "@hookform/resolvers/zod" import { useForm } from "react-hook-form" import * as z from "zod" import { Button } from "@/components/ui/button" import { Separator } from "@/components/ui/separator" import { Form, FormControl, FormDescription, FormField, FormItem, FormLabel, FormMessage, } from "@/components/ui/form" import { Input } from "@/components/ui/input" import { Switch } from "./ui/switch" import { Tabs, TabsContent, TabsList, TabsTrigger } from "./ui/tabs" import { useEffect, useState } from "react" import { cn } from "@/lib/utils" import { useQuery } from "@tanstack/react-query" import { fetchModelInfos, switchModel } from "@/lib/api" import { ModelInfo } from "@/lib/types" import { useStore } from "@/lib/states" import { ScrollArea } from "./ui/scroll-area" import { useToast } from "./ui/use-toast" import { AlertDialog, AlertDialogContent, AlertDialogHeader, } from "./ui/alert-dialog" import { MODEL_TYPE_DIFFUSERS_SD, MODEL_TYPE_DIFFUSERS_SDXL, MODEL_TYPE_DIFFUSERS_SDXL_INPAINT, MODEL_TYPE_DIFFUSERS_SD_INPAINT, MODEL_TYPE_INPAINT, MODEL_TYPE_OTHER, } from "@/lib/const" import useHotKey from "@/hooks/useHotkey" const formSchema = z.object({ enableFileManager: z.boolean(), inputDirectory: z.string().refine(async (id) => { // verify that ID exists in database return true }), outputDirectory: z.string().refine(async (id) => { // verify that ID exists in database return true }), enableDownloadMask: z.boolean(), enableManualInpainting: z.boolean(), enableUploadMask: z.boolean(), }) const TAB_GENERAL = "General" const TAB_MODEL = "Model" const TAB_FILE_MANAGER = "File Manager" const TAB_NAMES = [TAB_MODEL, TAB_GENERAL] export function SettingsDialog() { const [open, toggleOpen] = useToggle(false) const [openModelSwitching, toggleOpenModelSwitching] = useToggle(false) const [tab, setTab] = useState(TAB_MODEL) const [ updateAppState, settings, updateSettings, fileManagerState, updateFileManagerState, setAppModel, ] = useStore((state) => [ state.updateAppState, state.settings, state.updateSettings, state.fileManagerState, state.updateFileManagerState, state.setModel, ]) const { toast } = useToast() const [model, setModel] = useState(settings.model) useEffect(() => { setModel(settings.model) }, [settings.model]) const { data: modelInfos, status } = useQuery({ queryKey: ["modelInfos"], queryFn: fetchModelInfos, }) // 1. Define your form. const form = useForm>({ resolver: zodResolver(formSchema), defaultValues: { enableDownloadMask: settings.enableDownloadMask, enableManualInpainting: settings.enableManualInpainting, enableUploadMask: settings.enableUploadMask, inputDirectory: fileManagerState.inputDirectory, outputDirectory: fileManagerState.outputDirectory, }, }) function onSubmit(values: z.infer) { // Do something with the form values. ✅ This will be type-safe and validated. updateSettings({ enableDownloadMask: values.enableDownloadMask, enableManualInpainting: values.enableManualInpainting, enableUploadMask: values.enableUploadMask, }) // TODO: validate input/output Directory updateFileManagerState({ inputDirectory: values.inputDirectory, outputDirectory: values.outputDirectory, }) if (model.name !== settings.model.name) { toggleOpenModelSwitching() updateAppState({ disableShortCuts: true }) switchModel(model.name) .then((res) => { if (res.ok) { toast({ title: `Switch to ${model.name} success`, }) setAppModel(model) } else { throw new Error("Server error") } }) .catch(() => { toast({ variant: "destructive", title: `Switch to ${model.name} failed`, }) setModel(settings.model) }) .finally(() => { toggleOpenModelSwitching() updateAppState({ disableShortCuts: false }) }) } } useHotKey( "s", () => { toggleOpen() if (open) { onSubmit(form.getValues()) } }, [open, form, model] ) function onOpenChange(value: boolean) { toggleOpen() if (!value) { onSubmit(form.getValues()) } } function onModelSelect(info: ModelInfo) { setModel(info) } function renderModelList(model_types: string[]) { if (!modelInfos) { return
Please download model first
} return modelInfos .filter((info) => model_types.includes(info.model_type)) .map((info: ModelInfo) => { return (
onModelSelect(info)}>
{info.name}
) }) } function renderModelSettings() { if (status !== "success") { return <> } let defaultTab = MODEL_TYPE_INPAINT for (let info of modelInfos) { if (model.name === info.name) { defaultTab = info.model_type if (defaultTab === MODEL_TYPE_DIFFUSERS_SDXL) { defaultTab = MODEL_TYPE_DIFFUSERS_SD } if (defaultTab === MODEL_TYPE_DIFFUSERS_SDXL_INPAINT) { defaultTab = MODEL_TYPE_DIFFUSERS_SD_INPAINT } break } } return (
Current Model
{model.name}
Available models
{/* */}
Erase Stable Diffusion Stable Diffusion Inpaint Other Diffusion {renderModelList([MODEL_TYPE_INPAINT])} {renderModelList([ MODEL_TYPE_DIFFUSERS_SD, MODEL_TYPE_DIFFUSERS_SDXL, ])} {renderModelList([ MODEL_TYPE_DIFFUSERS_SD_INPAINT, MODEL_TYPE_DIFFUSERS_SDXL_INPAINT, ])} {renderModelList([MODEL_TYPE_OTHER])}
) } function renderGeneralSettings() { return (
(
Enable manual inpainting For erase model, click a button to trigger inpainting after draw mask.
)} /> (
Enable download mask Also download the mask after save the inpainting result.
)} /> {/* (
Enable upload mask Enable upload custom mask to perform inpainting.
)} /> */}
) } function renderFileManagerSettings() { return (
(
Enable file manger Browser images
)} /> ( Input directory Browser images from this directory. )} /> ( Save directory Result images will be saved to this directory. )} />
) } return ( <> {/* */}
Loading...
Switching to {model.name}
{/*
*/}
event.preventDefault()} onOpenAutoFocus={(event) => event.preventDefault()} // onPointerDownOutside={(event) => event.preventDefault()} > Settings
{TAB_NAMES.map((item) => ( ))}
{tab === TAB_MODEL ? renderModelSettings() : <>} {tab === TAB_GENERAL ? renderGeneralSettings() : <>} {/* {tab === TAB_FILE_MANAGER ? ( renderFileManagerSettings() ) : ( <> )} */}
) } export default SettingsDialog