diff --git a/src/components/modals/ModelSearchDialog.tsx b/src/components/modals/ModelSearchDialog.tsx index 29439cc5..acd57a1f 100644 --- a/src/components/modals/ModelSearchDialog.tsx +++ b/src/components/modals/ModelSearchDialog.tsx @@ -227,19 +227,46 @@ export function ModelSearchDialog({ } }; - // Get capability badge - const getCapabilityBadge = (capabilities: ModelCapability[]) => { - if ( - capabilities.includes("text-to-video") || - capabilities.includes("image-to-video") - ) { - return ( - - Video - - ); - } - return null; + // Get capability badges - show all capabilities to differentiate similar models + const getCapabilityBadges = (capabilities: ModelCapability[]) => { + const badges: React.ReactNode[] = []; + + capabilities.forEach((cap) => { + let color = ""; + let label = ""; + + switch (cap) { + case "text-to-image": + color = "bg-green-500/20 text-green-300"; + label = "txt→img"; + break; + case "image-to-image": + color = "bg-cyan-500/20 text-cyan-300"; + label = "img→img"; + break; + case "text-to-video": + color = "bg-purple-500/20 text-purple-300"; + label = "txt→vid"; + break; + case "image-to-video": + color = "bg-pink-500/20 text-pink-300"; + label = "img→vid"; + break; + } + + if (label) { + badges.push( + + {label} + + ); + } + }); + + return badges; }; if (!isOpen) return null; @@ -411,10 +438,10 @@ export function ModelSearchDialog({