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..2a6e955 --- /dev/null +++ b/ui/src/model-manager/install-models/AddApiKeyPopover.tsx @@ -0,0 +1,55 @@ +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 {} + +export default function AddApiKeyPopover({}: Props) { + const [apiKeyInput, setApiKeyInput] = useState(""); + const { onOpen, onClose, isOpen } = useDisclosure(); + + const saveApiKey = () => { + setCivitApiKey(apiKeyInput); + onClose(); + }; + + return ( + + + + + + + + + setApiKeyInput(e.target.value)} + placeholder="API Key" + /> + + + + + ); +} diff --git a/ui/src/model-manager/install-models/ChooseFolder.tsx b/ui/src/model-manager/install-models/ChooseFolder.tsx index ffd96aa..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 ece9828..6557ac9 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 { getCivitApiKey } from "../../utils/civitUtils"; type CivitModelQueryParams = { types?: MODEL_TYPE; @@ -143,11 +145,16 @@ 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}`; + } installModelsApi({ filename: file.current.name, name: file.current.name, save_path: folderPath, - url: file.current.downloadUrl, + url, }); onClose(); }; @@ -165,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(); }; @@ -194,6 +196,7 @@ export default function InatallModelsModal({ + {showInstallModels && ( - setShowInstallModels(false)} /> + setShowInstallModels(false)} + /> )} ); diff --git a/ui/src/model-manager/install-models/InstallProgress.tsx b/ui/src/model-manager/install-models/InstallProgress.tsx index 160d968..73d5e10 100644 --- a/ui/src/model-manager/install-models/InstallProgress.tsx +++ b/ui/src/model-manager/install-models/InstallProgress.tsx @@ -44,7 +44,7 @@ export default function InstallProgress() { paddingX={5} > {queue.map(({ save_path, progress }) => ( - + {save_path.replace(/^.*[\\/]/, "")} diff --git a/ui/src/model-manager/missing-models-drawer/MissingModelItem.tsx b/ui/src/model-manager/missing-models-drawer/MissingModelItem.tsx index 79df4f4..4afc565 100644 --- a/ui/src/model-manager/missing-models-drawer/MissingModelItem.tsx +++ b/ui/src/model-manager/missing-models-drawer/MissingModelItem.tsx @@ -15,22 +15,21 @@ interface Props { export default function MissingModelItem({ model }: Props) { const [suggestedUrls, setSuggestedUrls] = useState<{ url: string, name: string }[]>([]); const modelName = formatSearchQuery(model.received_value); - const getSearchData = async () => { - // huggingface search + const getHuggingFaceData = async () => { const hfData = await fetch(`https://huggingface.co/api/models?limit=3&search=${modelName}`); const hfSearchResult = await hfData.json() as { id: string; modelId: string }[]; const hfUrls = hfSearchResult.map(({ modelId }) => ({ name: `${modelId.split('/')[1]} on HuggingFace`, url: `https://huggingface.co/${modelId}` })); - console.log(modelName, "hf", hfSearchResult, hfUrls); - // civitai search + setSuggestedUrls(p => [...p, ...hfUrls]); + }; + const getCivitaiData = async () => { const civitaiData = await fetch(`https://civitai.com/api/v1/models?limit=3&query=${modelName}`); const civitaiSearchResult = await civitaiData.json() as { items: { id: string, name: string }[] }; const civitaiUrls = civitaiSearchResult.items?.map(({ name, id }) => ({ name: `${name} on civitAI`, url: `https://civitai.com/models/${id}` })) ?? []; - console.log(modelName, "civitai", civitaiSearchResult, civitaiUrls); - // set suggested urls - setSuggestedUrls([...hfUrls, ...civitaiUrls]); + setSuggestedUrls(p => [...p, ...civitaiUrls]); }; useEffect(() => { - getSearchData(); + getHuggingFaceData(); + getCivitaiData(); }, [model]); console.log(suggestedUrls); @@ -83,8 +82,8 @@ function formatSearchQuery(query: string): string { .replace(/^.*(\\|\/|\:)/, '') // Remove file extension .replace(/\.[^/.]+$/, "") - // Replace special characters with space - .replace(/[^a-zA-Z0-9]/g, " ") + // Replace underscore with space + .replaceAll("_", " ") // Add space before capital letters .replace(/([A-Z])/g, " $1") .trim(); diff --git a/ui/src/model-manager/models-list-drawer/ModelsList.tsx b/ui/src/model-manager/models-list-drawer/ModelsList.tsx index 318489a..dba2472 100644 --- a/ui/src/model-manager/models-list-drawer/ModelsList.tsx +++ b/ui/src/model-manager/models-list-drawer/ModelsList.tsx @@ -9,13 +9,11 @@ interface Props { export function ModelsList({ list }: Props) { return ( - {list.map((v) => { - return ( - - - - ); - })} + {list.map((v) => ( + + + + ))} ); } diff --git a/ui/src/model-manager/models-list-drawer/ModelsTags.tsx b/ui/src/model-manager/models-list-drawer/ModelsTags.tsx index 7220f68..5ebba8e 100644 --- a/ui/src/model-manager/models-list-drawer/ModelsTags.tsx +++ b/ui/src/model-manager/models-list-drawer/ModelsTags.tsx @@ -17,20 +17,18 @@ export function ModelsTags({ return ( - {modelTypeList.map((v) => { - return ( - - - - ); - })} + {modelTypeList.map((v) => ( + + + + ))} ); } 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); +}