Merge pull request #154 from 11cafe/feat/api-key-input

feat: api key input

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'<title>(.*?)</title>', 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))
 

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<Props>) {
   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,

ui/src/components/DropdownTitle.tsx

@@ -121,7 +121,7 @@ export default function DropdownTitle() {
   return (
     <>
       <Menu isLazy={true} isOpen={isOpen} gutter={4}>
-        <MenuButton onMouseEnter={onOpen} onMouseLeave={delayedClose}>
+        <MenuButton as="div" onMouseEnter={onOpen} onMouseLeave={delayedClose}>
           <IconButton
             icon={<IconChevronDown size={20} />}
             aria-label="menu"

ui/src/components/EditFlowName.tsx

@@ -76,6 +76,7 @@ export default function EditFlowName({
     <HStack>
       <Tooltip label={displayName} placement="bottom">
         <Text
+          as="div"
           color="white"
           onClick={startEdit}
           maxW={240}

ui/src/model-manager/hooks/useLocalStorage.ts

@@ -0,0 +1,41 @@
+import { useEffect, useState } from "react";
+
+export default function useLocalStorage<T>(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];
+}

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 (
+    <Popover
+      isOpen={isOpen}
+      onOpen={onOpen}
+      onClose={onClose}
+      placement="right"
+      closeOnBlur={false}
+    >
+      <PopoverTrigger>
+        <Button size={"sm"} py={1} mr={8}>
+          Set API Key
+        </Button>
+      </PopoverTrigger>
+      <PopoverContent p={5}>
+        <PopoverArrow />
+        <PopoverCloseButton />
+        <Stack spacing={4}>
+          <Input
+            value={apiKeyInput}
+            onChange={(e) => setApiKeyInput(e.target.value)}
+            placeholder="API Key"
+          />
+          <Button size={"sm"} py={1} mr={8} onClick={saveApiKey}>
+            Save
+          </Button>
+        </Stack>
+      </PopoverContent>
+    </Popover>
+  );
+}

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<string[]>([]);
-
+  const [url, setUrl] = useState<string>("");
   const cancelRef = useRef(null);
 
   useEffect(() => {
@@ -55,13 +57,21 @@ export default function ChooseFolder({
           </AlertDialogHeader>
 
           <AlertDialogBody>
+            <Text>Model download url</Text>
+            <Input
+              placeholder="https://civitai.com/api/download/models/311399"
+              onChange={(e) => setUrl(e.target.value)}
+              value={url}
+            />
             <Select
               placeholder="Select option"
               value={folderPath}
               onChange={(e) => setFolderPath(e.target.value)}
             >
               {foldersList.map((folderPath) => (
-                <option value={folderPath}>{folderPath}</option>
+                <option key={folderPath} value={folderPath}>
+                  {folderPath}
+                </option>
               ))}
             </Select>
           </AlertDialogBody>
@@ -70,7 +80,11 @@ export default function ChooseFolder({
             <Button ref={cancelRef} onClick={onClose}>
               Cancel
             </Button>
-            <Button onClick={() => selectFolder(folderPath)} ml={3}>
+            <Button
+              onClick={() => selectFolder(folderPath, url)}
+              ml={3}
+              isDisabled={url.length === 0}
+            >
               Confirm
             </Button>
           </AlertDialogFooter>

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({
               <Button size={"sm"} py={1} mr={8} onClick={customUrlDownload}>
                 Custom URL Install
               </Button>
+              <AddApiKeyPopover />
             </HStack>
             <HStack gap={2} mb={2} wrap={"wrap"}>
               <Button
@@ -209,6 +212,7 @@ export default function InatallModelsModal({
               {ALL_MODEL_TYPES.map((type) => {
                 return (
                   <Button
+                    key={type}
                     size={"sm"}
                     py={1}
                     isActive={modelType === type}
@@ -251,7 +255,10 @@ export default function InatallModelsModal({
       <ChooseFolder
         isOpen={isOpen}
         onClose={onClose}
-        selectFolder={downloadModels}
+        selectFolder={(folderPath: string, customUrl: string) => {
+          file.current = { id: 0, downloadUrl: customUrl };
+          downloadModels(folderPath);
+        }}
       />
     </>
   );

ui/src/model-manager/install-models/InstallModelsButton.tsx

@@ -1,5 +1,5 @@
-import { useRef, useState } from "react";
-import { Box, Button, Portal } from "@chakra-ui/react";
+import { useState } from "react";
+import { Button } from "@chakra-ui/react";
 
 import InatallModelsModal from "./InatallModelsModal";
 
@@ -17,7 +17,10 @@ export default function InstallModelsButton() {
       </Button>
 
       {showInstallModels && (
-        <InatallModelsModal modelType="Checkpoint" onclose={() => setShowInstallModels(false)} />
+        <InatallModelsModal
+          modelType="Checkpoint"
+          onclose={() => setShowInstallModels(false)}
+        />
       )}
     </>
   );

ui/src/model-manager/install-models/InstallProgress.tsx

@@ -44,7 +44,7 @@ export default function InstallProgress() {
       paddingX={5}
     >
       {queue.map(({ save_path, progress }) => (
-        <HStack>
+        <HStack key={save_path}>
           <Text fontSize={16} width="40%">
             {save_path.replace(/^.*[\\/]/, "")}
           </Text>

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();

ui/src/model-manager/models-list-drawer/ModelsList.tsx

@@ -9,13 +9,11 @@ interface Props {
 export function ModelsList({ list }: Props) {
   return (
     <Grid templateColumns="repeat(3, 1fr)" gap={1} marginTop={2}>
-      {list.map((v) => {
-        return (
-          <GridItem>
-            <ModelItem data={v} key={v.model_name} />
-          </GridItem>
-        );
-      })}
+      {list.map((v) => (
+        <GridItem key={v.model_name}>
+          <ModelItem data={v} />
+        </GridItem>
+      ))}
     </Grid>
   );
 }

ui/src/model-manager/models-list-drawer/ModelsTags.tsx

@@ -17,20 +17,18 @@ export function ModelsTags({
 
   return (
     <Wrap>
-      {modelTypeList.map((v) => {
-        return (
-          <WrapItem>
-            <Button
-              colorScheme="blue"
-              variant={selectedModel === v ? "solid" : "outline"}
-              onClick={() => clickHanlder(v)}
-              size={"sm"}
-            >
-              {v}
-            </Button>
-          </WrapItem>
-        );
-      })}
+      {modelTypeList.map((v) => (
+        <WrapItem key={v}>
+          <Button
+            colorScheme="blue"
+            variant={selectedModel === v ? "solid" : "outline"}
+            onClick={() => clickHanlder(v)}
+            size={"sm"}
+          >
+            {v}
+          </Button>
+        </WrapItem>
+      ))}
     </Wrap>
   );
 }

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);
+}