feat: missing model suggestion

2024-01-25 11:26:58 +08:00
parent 44d040fadd
commit b23b3c7c33
3 changed files with 121 additions and 89 deletions
@@ -0,0 +1,89 @@
+import { GridItem, VStack, Button, Text } from "@chakra-ui/react";
+import { IconExternalLink } from "@tabler/icons-react";
+import { useEffect, useState } from "react";
+
+export interface MissingModel {
+  class_type: string;
+  input_name: string;
+  received_value: string;
+}
+
+interface Props {
+  model: MissingModel;
+}
+
+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 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
+    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]);
+  };
+  useEffect(() => {
+    getSearchData();
+  }, [model]);
+
+  console.log(suggestedUrls);
+
+  return (
+    <GridItem p={3} shadow="md" borderWidth="1px">
+      <VStack align="start">
+        <Text>{model.class_type}</Text>
+        <Text fontWeight="bold">Input:</Text>
+        <Text>{model.input_name}</Text>
+        <Text fontWeight="bold">Received Value:</Text>
+        <Text color="red.400">{model.received_value}</Text>
+        {suggestedUrls.map(({ url, name }) => (
+          <Button
+            colorScheme="blue"
+            mt={5}
+            iconSpacing={1}
+            leftIcon={<IconExternalLink size={20} />}
+            size={"sm"}
+            onClick={() => {
+              window.open(url, "_blank");
+            }}
+          >
+            {name}
+          </Button>
+        ))}
+        <Button
+          colorScheme="blue"
+          mt={5}
+          iconSpacing={1}
+          leftIcon={<IconExternalLink size={20} />}
+          size={"sm"}
+          onClick={() => {
+            window.open(
+              `https://civitai.com/search/models?sortBy=models_v5&query=${modelName}`,
+              "_blank",
+            );
+          }}
+        >
+          Search in CivitAI
+        </Button>
+      </VStack>
+    </GridItem>
+  );
+}
+
+function formatSearchQuery(query: string): string {
+  return query
+    // Remove file extension
+    .replace(/\.[^/.]+$/, "")
+    // Replace special characters with space
+    .replace(/[^a-zA-Z0-9]/g, " ")
+    // Add space before capital letters
+    .replace(/([A-Z])/g, " $1")
+    .trim();
+}
@@ -1,22 +1,17 @@
 import {
  Box,
-  Button,
  Card,
  CardHeader,
  Flex,
  Grid,
-  GridItem,
  Heading,
  Portal,
-  Text,
-  VStack,
 } from "@chakra-ui/react";
-import { useEffect, useState } from "react";
+import { useEffect } from "react";
 // @ts-ignore
 import { app } from "/scripts/app.js";
 import InstallModelsButton from "../install-models/InstallModelsButton";
-import InatallModelsModal from "../install-models/InatallModelsModal";
-import { IconExternalLink } from "@tabler/icons-react";
+import MissingModelItem from "./MissingModelItem";

 export interface MissingModel {
  class_type: string;
@@ -33,9 +28,6 @@ export default function MissingModelsListDrawer({
  onClose,
  missingModels,
 }: Props) {
-  const [showInstallModels, setShowInstallModels] = useState(false);
-  const [searchQuery, setSearchQuery] = useState("");
-
  useEffect(() => {
    app.canvasEl.addEventListener("click", onClose);
    return () => {
@@ -46,81 +38,33 @@ export default function MissingModelsListDrawer({
  const DRAWER_WIDTH = 440;

  return (
-    <>
-      <Portal>
-        <Box style={{ width: DRAWER_WIDTH }}>
-          <Card
-            width={DRAWER_WIDTH}
-            height={"100vh"}
-            p={4}
-            gap={2}
-            position={"fixed"}
-            top={0}
-            left={0}
-            shadow={"xl"}
-            zIndex={1000}
-            overflowY={"auto"}
-          >
-            <CardHeader>
-              <Flex justifyContent={"space-between"} alignContent={"center"}>
-                <Heading size={"md"} mr={2}>
-                  Models
-                </Heading>
-                <InstallModelsButton />
-              </Flex>
-            </CardHeader>
-            <Grid templateColumns="1" gap={1} marginTop={2}>
-              {missingModels.map((model) => {
-                return (
-                  <GridItem p={3} shadow="md" borderWidth="1px">
-                    <VStack align="start">
-                      <Text>{model.class_type}</Text>
-                      <Text fontWeight="bold">Input:</Text>
-                      <Text>{model.input_name}</Text>
-                      <Text fontWeight="bold">Received Value:</Text>
-                      <Text color="red.400">{model.received_value}</Text>
-                    </VStack>
-                    <Button
-                      colorScheme="blue"
-                      mt={5}
-                      iconSpacing={1}
-                      leftIcon={<IconExternalLink size={20} />}
-                      size={"sm"}
-                      onClick={() => {
-                        window.open(
-                          `https://civitai.com/search/models?sortBy=models_v5&query=${formatSearchQuery(model.received_value)}`,
-                          "_blank",
-                        );
-                      }}
-                    >
-                      Search in CivitAI
-                    </Button>
-                  </GridItem>
-                );
-              })}
-            </Grid>
-          </Card>
-        </Box>
-      </Portal>
-      {showInstallModels && (
-        <InatallModelsModal
-          searchQuery={searchQuery}
-          onclose={() => setShowInstallModels(false)}
-        />
-      )}
-    </>
+    <Portal>
+      <Box style={{ width: DRAWER_WIDTH }}>
+        <Card
+          width={DRAWER_WIDTH}
+          height={"100vh"}
+          p={4}
+          gap={2}
+          position={"fixed"}
+          top={0}
+          left={0}
+          shadow={"xl"}
+          zIndex={1000}
+          overflowY={"auto"}
+        >
+          <CardHeader>
+            <Flex justifyContent={"space-between"} alignContent={"center"}>
+              <Heading size={"md"} mr={2}>
+                Models
+              </Heading>
+              <InstallModelsButton />
+            </Flex>
+          </CardHeader>
+          <Grid templateColumns="1" gap={1} marginTop={2}>
+            {missingModels.map((model) => <MissingModelItem model={model} />)}
+          </Grid>
+        </Card>
+      </Box>
+    </Portal>
  );
 }
-
-function formatSearchQuery(query: string): string {
-  // Remove file extension
-  let formattedQuery = query
-    .replace(/\.[^/.]+$/, "")
-    // Replace special characters with space
-    // .replace(/[^a-zA-Z0-9]/g, " ")
-    // Add space before capital letters
-    // .replace(/([A-Z])/g, " $1")
-    .trim();
-
-  return formattedQuery;
-}
@@ -1,11 +1,10 @@
 import { Box, Button, HStack } from "@chakra-ui/react";
-import { useEffect, useState, DragEvent } from "react";
-import React from "react";
+import { lazy, useEffect, useState, DragEvent } from "react";
 import ModelsListDrawer from "../models-list-drawer/ModelsListDrawer";
 // @ts-ignore
 import { app } from "/scripts/app.js";

-const AddMissingModelsButton = React.lazy(
+const AddMissingModelsButton = lazy(
  () => import("./InstallMissingModelsButton"),
 );
 interface Props {}