diff --git a/ui/src/model-manager/missing-models-drawer/MissingModelItem.tsx b/ui/src/model-manager/missing-models-drawer/MissingModelItem.tsx index 4afc565..7c0663e 100644 --- a/ui/src/model-manager/missing-models-drawer/MissingModelItem.tsx +++ b/ui/src/model-manager/missing-models-drawer/MissingModelItem.tsx @@ -1,7 +1,7 @@ -import { GridItem, VStack, Button, Text } from "@chakra-ui/react"; +import { GridItem, VStack, Button, Text, Link } from "@chakra-ui/react"; import { IconExternalLink } from "@tabler/icons-react"; import { useEffect, useState } from "react"; - +import { modelExtensions } from "../utils"; export interface MissingModel { class_type: string; input_name: string; @@ -11,28 +11,57 @@ export interface MissingModel { interface Props { model: MissingModel; } - +type missingModelPredict = { url: string; name: string; downloadUrl: string }; export default function MissingModelItem({ model }: Props) { - const [suggestedUrls, setSuggestedUrls] = useState<{ url: string, name: string }[]>([]); + const [suggestedUrls, setSuggestedUrls] = useState([]); const modelName = formatSearchQuery(model.received_value); 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}` })); - setSuggestedUrls(p => [...p, ...hfUrls]); + console.warn("fetching hf", modelName); + const hfData = await fetch( + `https://huggingface.co/api/models?limit=3&search=${modelName}&full=true`, + ); + const hfSearchResult = (await hfData.json()) as { + id: string; + modelId: string; + siblings?: { rfilename: string }[]; + }[]; + const hfUrls: missingModelPredict[] = []; + console.log(hfSearchResult); + hfSearchResult.slice(0, 1).forEach(({ modelId, siblings }) => { + console.log(modelId, siblings); + siblings?.forEach((s) => { + const ext = s.rfilename.split(".").pop(); + console.log(ext); + if (modelExtensions.includes(ext ?? "")) { + hfUrls.push({ + name: s.rfilename, + downloadUrl: `https://huggingface.co/${modelId}/resolve/main/${s.rfilename}`, + url: `https://huggingface.co/${modelId}/tree/main`, + }); + } + }); + }); + setSuggestedUrls(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}` })) ?? []; - setSuggestedUrls(p => [...p, ...civitaiUrls]); + 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}`, + downloadUrl: `https://civitai.com/models/${id}`, + })) ?? []; + setSuggestedUrls((p) => [...p, ...civitaiUrls]); }; useEffect(() => { getHuggingFaceData(); - getCivitaiData(); - }, [model]); - - console.log(suggestedUrls); + // getCivitaiData(); + }, []); return ( @@ -43,18 +72,16 @@ export default function MissingModelItem({ model }: Props) { Received Value: {model.received_value} {suggestedUrls.map(({ url, name }) => ( - + + ))}