From 96c3f3385cfb8637145c9f2aef38eb1cb1bb9723 Mon Sep 17 00:00:00 2001 From: Arslan Ablikim Date: Sat, 27 Jan 2024 13:47:51 +0800 Subject: [PATCH 1/2] feat: api key input --- service/model_manager/model_installer.py | 12 ++++++ ui/src/components/Draggable.tsx | 15 +++++-- ui/src/components/DropdownTitle.tsx | 2 +- ui/src/components/EditFlowName.tsx | 1 + ui/src/model-manager/hooks/useLocalStorage.ts | 41 ++++++++++++++++++ .../install-models/AddApiKeyPopover.tsx | 43 +++++++++++++++++++ .../install-models/ChooseFolder.tsx | 2 +- .../install-models/InatallModelsModal.tsx | 13 +++++- .../install-models/InstallProgress.tsx | 2 +- .../MissingModelItem.tsx | 19 ++++---- .../models-list-drawer/ModelsList.tsx | 12 +++--- .../models-list-drawer/ModelsTags.tsx | 26 ++++++----- 12 files changed, 148 insertions(+), 40 deletions(-) create mode 100644 ui/src/model-manager/hooks/useLocalStorage.ts create mode 100644 ui/src/model-manager/install-models/AddApiKeyPopover.tsx diff --git a/service/model_manager/model_installer.py b/service/model_manager/model_installer.py index 5c45351..4028fe3 100644 --- a/service/model_manager/model_installer.py +++ b/service/model_manager/model_installer.py @@ -147,10 +147,22 @@ def download_url_with_agent(url, save_path, progress_callback=None): req = urllib.request.Request(url, headers=headers) with urllib.request.urlopen(req) as response: + if response.status != 200: + raise Exception(f"Request returned status code {response.status}") + file_size = int(response.headers.get('content-length', 0)) + print(f"File size: {file_size} bytes") chunk_size = 1024 # 1KB per chunk downloaded = 0 + if file_size == 0: + html = response.read().decode() + title = re.search(r'(.*?)', html)[0] + if 'Sign in' in title: + raise Exception(f"You need to add API key to download this model.") + else: + raise Exception(f"Download failed. {title}") + if not os.path.exists(os.path.dirname(temp_save_path)): os.makedirs(os.path.dirname(temp_save_path)) diff --git a/ui/src/components/Draggable.tsx b/ui/src/components/Draggable.tsx index 9ae4b04..0850aac 100644 --- a/ui/src/components/Draggable.tsx +++ b/ui/src/components/Draggable.tsx @@ -1,17 +1,24 @@ -import { useState, useCallback, MouseEvent, useRef } from "react"; +import { useState, useCallback, MouseEvent, useRef, PropsWithChildren } from "react"; const POSITION = { x: 0, y: 0 }; -export default function Draggable({ children, onDragEnd }) { +interface Props { + onDragEnd: (position: { x: number; y: number }) => void; +} + +export default function Draggable({ children, onDragEnd }: PropsWithChildren) { const originRef = useRef(POSITION); const isDraggingRef = useRef(false); const [translation, setTranslation] = useState(POSITION); const handleMouseDown = (e: MouseEvent) => { + if (!(e.target instanceof Element) || !(e.target.parentNode instanceof Element)) { + return; + } if ([e.target?.id, e.target?.parentNode?.id].includes("dragPanelIcon")) { originRef.current = { x: e.clientX, - y: e.clientY, + y: e.clientY, }; isDraggingRef.current = true; window.addEventListener("mousemove", handleMouseMove); @@ -21,7 +28,7 @@ export default function Draggable({ children, onDragEnd }) { } }; - const handleMouseMove = useCallback((e: MouseEvent) => { + const handleMouseMove = useCallback((e: globalThis.MouseEvent) => { const translation = { x: e.clientX - originRef.current.x, y: e.clientY - originRef.current.y, diff --git a/ui/src/components/DropdownTitle.tsx b/ui/src/components/DropdownTitle.tsx index 8e15a6f..3ba2f9b 100644 --- a/ui/src/components/DropdownTitle.tsx +++ b/ui/src/components/DropdownTitle.tsx @@ -121,7 +121,7 @@ export default function DropdownTitle() { return ( <> - + } aria-label="menu" diff --git a/ui/src/components/EditFlowName.tsx b/ui/src/components/EditFlowName.tsx index 70bd303..e4a5fe0 100644 --- a/ui/src/components/EditFlowName.tsx +++ b/ui/src/components/EditFlowName.tsx @@ -76,6 +76,7 @@ export default function EditFlowName({ (key: string, defaultValue: T): [T, (value: T) => void] { + const [value, setValue] = useState(defaultValue); + + useEffect(() => { + const item = localStorage.getItem(key); + + if (!item) { + localStorage.setItem(key, JSON.stringify(defaultValue)) + } + + setValue(item ? JSON.parse(item) : defaultValue) + + function handler(e: StorageEvent) { + if (e.key !== key) return; + + const lsi = localStorage.getItem(key) + setValue(JSON.parse(lsi ?? "")) + } + + window.addEventListener("storage", handler) + + return () => { + window.removeEventListener("storage", handler) + }; + }, [defaultValue, key]) + + const setValueWrap = (value: T) => { + try { + setValue(value); + + localStorage.setItem(key, JSON.stringify(value)); + if (typeof window !== "undefined") { + window.dispatchEvent(new StorageEvent("storage", { key })) + } + } catch (e) { console.error(e) } + }; + + return [value, setValueWrap]; +} \ No newline at end of file diff --git a/ui/src/model-manager/install-models/AddApiKeyPopover.tsx b/ui/src/model-manager/install-models/AddApiKeyPopover.tsx new file mode 100644 index 0000000..4b2ecbb --- /dev/null +++ b/ui/src/model-manager/install-models/AddApiKeyPopover.tsx @@ -0,0 +1,43 @@ +import { Popover, PopoverTrigger, PopoverContent, PopoverArrow, PopoverCloseButton, useDisclosure, Button, Input, Stack } from "@chakra-ui/react"; +import { useState } from "react"; + +interface Props { + apiKey: string; + setApiKey(apiKey: string): void; +} + +export default function AddApiKeyPopover({ apiKey, setApiKey }: Props) { + const [apiKeyInput, setApiKeyInput] = useState(apiKey); + const { onOpen, onClose, isOpen } = useDisclosure(); + + const saveApiKey = () => { + setApiKey(apiKeyInput); + onClose(); + } + + return ( + + + + + + + + + setApiKeyInput(e.target.value)} placeholder="API Key" /> + + + + + ); +} \ No newline at end of file diff --git a/ui/src/model-manager/install-models/ChooseFolder.tsx b/ui/src/model-manager/install-models/ChooseFolder.tsx index ffd96aa..8593332 100644 --- a/ui/src/model-manager/install-models/ChooseFolder.tsx +++ b/ui/src/model-manager/install-models/ChooseFolder.tsx @@ -61,7 +61,7 @@ export default function ChooseFolder({ onChange={(e) => setFolderPath(e.target.value)} > {foldersList.map((folderPath) => ( - + ))} diff --git a/ui/src/model-manager/install-models/InatallModelsModal.tsx b/ui/src/model-manager/install-models/InatallModelsModal.tsx index ece9828..7ae1d8c 100644 --- a/ui/src/model-manager/install-models/InatallModelsModal.tsx +++ b/ui/src/model-manager/install-models/InatallModelsModal.tsx @@ -20,6 +20,8 @@ import InstallModelSearchBar from "./InstallModelSearchBar"; import ChooseFolder from "./ChooseFolder"; import InstallProgress from "./InstallProgress"; import { indexdb } from "../../db-tables/indexdb"; +import AddApiKeyPopover from "./AddApiKeyPopover"; +import useLocalStorage from "../hooks/useLocalStorage"; type CivitModelQueryParams = { types?: MODEL_TYPE; @@ -74,6 +76,7 @@ export default function InatallModelsModal({ const [searchQuery, setSearchQuery] = useState(searchQueryProp); const { isOpen, onOpen, onClose } = useDisclosure(); const file = useRef(); + const [apiKey, setApiKey] = useLocalStorage('apiKey', ""); const loadData = useCallback(async () => { setLoading(true); const params: CivitModelQueryParams = { @@ -143,11 +146,15 @@ export default function InatallModelsModal({ }); file.current.name != null && setInstalling((cur) => [...cur, file.current?.name ?? ""]); + let url = file.current.downloadUrl; + if (apiKey) { + url += `?token=${apiKey}`; + } installModelsApi({ filename: file.current.name, name: file.current.name, save_path: folderPath, - url: file.current.downloadUrl, + url, }); onClose(); }; @@ -191,6 +198,7 @@ export default function InatallModelsModal({ setSearchQuery={setSearchQuery} onSearch={loadData} /> + @@ -209,6 +217,7 @@ export default function InatallModelsModal({ {ALL_MODEL_TYPES.map((type) => { return ( - - ); - })} + {modelTypeList.map((v) => ( + + + + ))} ); } From 90958fb5df307c8dc15a39f33fcb7f07a4d36ac3 Mon Sep 17 00:00:00 2001 From: Weixuan Fu Date: Sat, 27 Jan 2024 15:13:30 +0800 Subject: [PATCH 2/2] ui fix, remove listener, put custom url input into dialog --- .../install-models/AddApiKeyPopover.tsx | 36 ++++++++++++------- .../install-models/ChooseFolder.tsx | 22 +++++++++--- .../install-models/InatallModelsModal.tsx | 18 +++++----- .../install-models/InstallModelsButton.tsx | 9 +++-- ui/src/utils/civitUtils.ts | 8 +++++ 5 files changed, 64 insertions(+), 29 deletions(-) create mode 100644 ui/src/utils/civitUtils.ts diff --git a/ui/src/model-manager/install-models/AddApiKeyPopover.tsx b/ui/src/model-manager/install-models/AddApiKeyPopover.tsx index 4b2ecbb..2a6e955 100644 --- a/ui/src/model-manager/install-models/AddApiKeyPopover.tsx +++ b/ui/src/model-manager/install-models/AddApiKeyPopover.tsx @@ -1,26 +1,34 @@ -import { Popover, PopoverTrigger, PopoverContent, PopoverArrow, PopoverCloseButton, useDisclosure, Button, Input, Stack } from "@chakra-ui/react"; +import { + Popover, + PopoverTrigger, + PopoverContent, + PopoverArrow, + PopoverCloseButton, + useDisclosure, + Button, + Input, + Stack, +} from "@chakra-ui/react"; import { useState } from "react"; +import { setCivitApiKey } from "../../utils/civitUtils"; -interface Props { - apiKey: string; - setApiKey(apiKey: string): void; -} +interface Props {} -export default function AddApiKeyPopover({ apiKey, setApiKey }: Props) { - const [apiKeyInput, setApiKeyInput] = useState(apiKey); +export default function AddApiKeyPopover({}: Props) { + const [apiKeyInput, setApiKeyInput] = useState(""); const { onOpen, onClose, isOpen } = useDisclosure(); const saveApiKey = () => { - setApiKey(apiKeyInput); + setCivitApiKey(apiKeyInput); onClose(); - } + }; return ( @@ -32,7 +40,11 @@ export default function AddApiKeyPopover({ apiKey, setApiKey }: Props) { - setApiKeyInput(e.target.value)} placeholder="API Key" /> + setApiKeyInput(e.target.value)} + placeholder="API Key" + /> @@ -40,4 +52,4 @@ export default function AddApiKeyPopover({ apiKey, setApiKey }: Props) { ); -} \ No newline at end of file +} diff --git a/ui/src/model-manager/install-models/ChooseFolder.tsx b/ui/src/model-manager/install-models/ChooseFolder.tsx index 8593332..57aaf04 100644 --- a/ui/src/model-manager/install-models/ChooseFolder.tsx +++ b/ui/src/model-manager/install-models/ChooseFolder.tsx @@ -7,6 +7,8 @@ import { AlertDialogBody, AlertDialogFooter, Select, + Text, + Input, } from "@chakra-ui/react"; import { useEffect, useRef, useState } from "react"; import { getAllFoldersList } from "../../Api"; @@ -14,7 +16,7 @@ import { getAllFoldersList } from "../../Api"; interface ChooseFolderProps { isOpen: boolean; onClose: () => void; - selectFolder: (folderPath: string) => void; + selectFolder: (folderPath: string, url: string) => void; } export default function ChooseFolder({ isOpen, @@ -23,7 +25,7 @@ export default function ChooseFolder({ }: ChooseFolderProps) { const [folderPath, setFolderPath] = useState(""); const [foldersList, setFoldersList] = useState([]); - + const [url, setUrl] = useState(""); const cancelRef = useRef(null); useEffect(() => { @@ -55,13 +57,21 @@ export default function ChooseFolder({ + Model download url + setUrl(e.target.value)} + value={url} + /> @@ -70,7 +80,11 @@ export default function ChooseFolder({ - diff --git a/ui/src/model-manager/install-models/InatallModelsModal.tsx b/ui/src/model-manager/install-models/InatallModelsModal.tsx index 7ae1d8c..6557ac9 100644 --- a/ui/src/model-manager/install-models/InatallModelsModal.tsx +++ b/ui/src/model-manager/install-models/InatallModelsModal.tsx @@ -21,7 +21,7 @@ import ChooseFolder from "./ChooseFolder"; import InstallProgress from "./InstallProgress"; import { indexdb } from "../../db-tables/indexdb"; import AddApiKeyPopover from "./AddApiKeyPopover"; -import useLocalStorage from "../hooks/useLocalStorage"; +import { getCivitApiKey } from "../../utils/civitUtils"; type CivitModelQueryParams = { types?: MODEL_TYPE; @@ -76,7 +76,6 @@ export default function InatallModelsModal({ const [searchQuery, setSearchQuery] = useState(searchQueryProp); const { isOpen, onOpen, onClose } = useDisclosure(); const file = useRef(); - const [apiKey, setApiKey] = useLocalStorage('apiKey', ""); const loadData = useCallback(async () => { setLoading(true); const params: CivitModelQueryParams = { @@ -147,6 +146,7 @@ export default function InatallModelsModal({ file.current.name != null && setInstalling((cur) => [...cur, file.current?.name ?? ""]); let url = file.current.downloadUrl; + const apiKey = getCivitApiKey(); if (apiKey) { url += `?token=${apiKey}`; } @@ -172,11 +172,6 @@ export default function InatallModelsModal({ } }; const customUrlDownload = () => { - const downloadUrl = prompt("Enter the URL to download"); - if (!downloadUrl) { - return; - } - file.current = { id: 0, downloadUrl }; onOpen(); }; @@ -198,10 +193,10 @@ export default function InatallModelsModal({ setSearchQuery={setSearchQuery} onSearch={loadData} /> - + {showInstallModels && ( - setShowInstallModels(false)} /> + setShowInstallModels(false)} + /> )} ); diff --git a/ui/src/utils/civitUtils.ts b/ui/src/utils/civitUtils.ts new file mode 100644 index 0000000..7f36efc --- /dev/null +++ b/ui/src/utils/civitUtils.ts @@ -0,0 +1,8 @@ +const CIVIT_API_KEY_STORAGE_KEY = "WORKSPACE_CIVIT_API_KEY_STORAGE_KEY"; +export function getCivitApiKey() { + return localStorage.getItem(CIVIT_API_KEY_STORAGE_KEY); +} + +export function setCivitApiKey(apiKey: string) { + localStorage.setItem(CIVIT_API_KEY_STORAGE_KEY, apiKey); +}