chore: import upstream snapshot with attribution

invokeai/backend/model_manager/util/libc_util.py

@@ -0,0 +1,78 @@
+import ctypes
+
+
+class Struct_mallinfo2(ctypes.Structure):
+    """A ctypes Structure that matches the libc mallinfo2 struct.
+
+    Docs:
+    - https://man7.org/linux/man-pages/man3/mallinfo.3.html
+    - https://www.gnu.org/software/libc/manual/html_node/Statistics-of-Malloc.html
+
+    struct mallinfo2 {
+        size_t arena;     /* Non-mmapped space allocated (bytes) */
+        size_t ordblks;   /* Number of free chunks */
+        size_t smblks;    /* Number of free fastbin blocks */
+        size_t hblks;     /* Number of mmapped regions */
+        size_t hblkhd;    /* Space allocated in mmapped regions (bytes) */
+        size_t usmblks;   /* See below */
+        size_t fsmblks;   /* Space in freed fastbin blocks (bytes) */
+        size_t uordblks;  /* Total allocated space (bytes) */
+        size_t fordblks;  /* Total free space (bytes) */
+        size_t keepcost;  /* Top-most, releasable space (bytes) */
+    };
+    """
+
+    _fields_ = [
+        ("arena", ctypes.c_size_t),
+        ("ordblks", ctypes.c_size_t),
+        ("smblks", ctypes.c_size_t),
+        ("hblks", ctypes.c_size_t),
+        ("hblkhd", ctypes.c_size_t),
+        ("usmblks", ctypes.c_size_t),
+        ("fsmblks", ctypes.c_size_t),
+        ("uordblks", ctypes.c_size_t),
+        ("fordblks", ctypes.c_size_t),
+        ("keepcost", ctypes.c_size_t),
+    ]
+
+    def __str__(self) -> str:
+        s = ""
+        s += (
+            f"{'arena': <10}= {(self.arena / 2**30):15.5f}   # Non-mmapped space allocated (GB) (uordblks + fordblks)\n"
+        )
+        s += f"{'ordblks': <10}= {(self.ordblks): >15}   # Number of free chunks\n"
+        s += f"{'smblks': <10}= {(self.smblks): >15}   # Number of free fastbin blocks \n"
+        s += f"{'hblks': <10}= {(self.hblks): >15}   # Number of mmapped regions \n"
+        s += f"{'hblkhd': <10}= {(self.hblkhd / 2**30):15.5f}   # Space allocated in mmapped regions (GB)\n"
+        s += f"{'usmblks': <10}= {(self.usmblks): >15}   # Unused\n"
+        s += f"{'fsmblks': <10}= {(self.fsmblks / 2**30):15.5f}   # Space in freed fastbin blocks (GB)\n"
+        s += (
+            f"{'uordblks': <10}= {(self.uordblks / 2**30):15.5f}   # Space used by in-use allocations (non-mmapped)"
+            " (GB)\n"
+        )
+        s += f"{'fordblks': <10}= {(self.fordblks / 2**30):15.5f}   # Space in free blocks (non-mmapped) (GB)\n"
+        s += f"{'keepcost': <10}= {(self.keepcost / 2**30):15.5f}   # Top-most, releasable space (GB)\n"
+        return s
+
+
+class LibcUtil:
+    """A utility class for interacting with the C Standard Library (`libc`) via ctypes.
+
+    Note that this class will raise on __init__() if 'libc.so.6' can't be found. Take care to handle environments where
+    this shared library is not available.
+
+    TODO: Improve cross-OS compatibility of this class.
+    """
+
+    def __init__(self) -> None:
+        self._libc = ctypes.cdll.LoadLibrary("libc.so.6")
+
+    def mallinfo2(self) -> Struct_mallinfo2:
+        """Calls `libc` `mallinfo2`.
+
+        Docs: https://man7.org/linux/man-pages/man3/mallinfo.3.html
+        """
+        mallinfo2 = self._libc.mallinfo2
+        mallinfo2.restype = Struct_mallinfo2
+        result: Struct_mallinfo2 = mallinfo2()
+        return result

invokeai/backend/model_manager/util/lora_metadata_extractor.py

@@ -0,0 +1,146 @@
+"""Utility functions for extracting metadata from LoRA model files."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional, Set, Tuple
+
+from PIL import Image
+
+from invokeai.app.util.thumbnails import make_thumbnail
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.taxonomy import ModelType
+
+logger = logging.getLogger(__name__)
+
+
+def extract_lora_metadata(
+    model_path: Path, model_key: str, model_images_path: Path
+) -> Tuple[Optional[str], Optional[Set[str]]]:
+    """
+    Extract metadata for a LoRA model from associated JSON and image files.
+
+    Args:
+        model_path: Path to the LoRA model file
+        model_key: Unique key for the model
+        model_images_path: Path to the model images directory
+
+    Returns:
+        Tuple of (description, trigger_phrases)
+    """
+    model_stem = model_path.stem
+    model_dir = model_path.parent
+
+    # Find and process preview image
+    _process_preview_image(model_stem, model_dir, model_key, model_images_path)
+
+    # Extract metadata from JSON
+    description, trigger_phrases = _extract_json_metadata(model_stem, model_dir)
+
+    return description, trigger_phrases
+
+
+def _process_preview_image(model_stem: str, model_dir: Path, model_key: str, model_images_path: Path) -> bool:
+    """Find and process a preview image for the model, saving it to the model images store."""
+    image_extensions = [".png", ".jpg", ".jpeg", ".webp"]
+
+    for ext in image_extensions:
+        image_path = model_dir / f"{model_stem}{ext}"
+        if image_path.exists():
+            try:
+                # Open the image
+                with Image.open(image_path) as img:
+                    # Create thumbnail and save to model images directory
+                    thumbnail = make_thumbnail(img, 256)
+                    thumbnail_path = model_images_path / f"{model_key}.webp"
+                    thumbnail.save(thumbnail_path, format="webp")
+
+                logger.info(f"Processed preview image {image_path.name} for model {model_key}")
+                return True
+
+            except Exception as e:
+                logger.warning(f"Failed to process preview image {image_path.name}: {e}")
+                return False
+
+    return False
+
+
+def _extract_json_metadata(model_stem: str, model_dir: Path) -> Tuple[Optional[str], Optional[Set[str]]]:
+    """Extract metadata from a JSON file with the same name as the model."""
+    json_path = model_dir / f"{model_stem}.json"
+
+    if not json_path.exists():
+        return None, None
+
+    try:
+        with open(json_path, "r", encoding="utf-8") as f:
+            metadata = json.load(f)
+
+        # Extract description
+        description = _build_description(metadata)
+
+        # Extract trigger phrases
+        trigger_phrases = _extract_trigger_phrases(metadata)
+
+        if description or trigger_phrases:
+            logger.info(f"Applied metadata from {json_path.name}")
+
+        return description, trigger_phrases
+
+    except (json.JSONDecodeError, IOError, Exception) as e:
+        logger.warning(f"Failed to read metadata from {json_path}: {e}")
+        return None, None
+
+
+def _build_description(metadata: Dict[str, Any]) -> Optional[str]:
+    """Build a description from metadata fields."""
+    description_parts = []
+
+    if description := metadata.get("description"):
+        description_parts.append(str(description).strip())
+
+    if notes := metadata.get("notes"):
+        description_parts.append(str(notes).strip())
+
+    return " | ".join(description_parts) if description_parts else None
+
+
+def _extract_trigger_phrases(metadata: Dict[str, Any]) -> Optional[Set[str]]:
+    """Extract trigger phrases from metadata."""
+    if not (activation_text := metadata.get("activation text")):
+        return None
+
+    activation_text = str(activation_text).strip()
+    if not activation_text:
+        return None
+
+    # Split on commas and clean up each phrase
+    phrases = [phrase.strip() for phrase in activation_text.split(",") if phrase.strip()]
+
+    return set(phrases) if phrases else None
+
+
+def apply_lora_metadata(info: AnyModelConfig, model_path: Path, model_images_path: Path) -> None:
+    """
+    Apply extracted metadata to a LoRA model configuration.
+
+    Args:
+        info: The model configuration to update
+        model_path: Path to the LoRA model file
+        model_images_path: Path to the model images directory
+    """
+    # Only process LoRA models
+    if info.type != ModelType.LoRA:
+        return
+
+    # Extract and apply metadata
+    description, trigger_phrases = extract_lora_metadata(model_path, info.key, model_images_path)
+
+    # We don't set cover_image path in the config anymore since images are stored
+    # separately in the model images store by model key
+
+    if description:
+        info.description = description
+
+    if trigger_phrases:
+        info.trigger_phrases = trigger_phrases

invokeai/backend/model_manager/util/model_util.py

@@ -0,0 +1,210 @@
+"""Utilities for parsing model files, used mostly by legacy_probe.py"""
+
+import json
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+import picklescan.scanner as pscan
+import safetensors
+import torch
+
+from invokeai.app.services.config.config_default import get_config
+from invokeai.backend.model_manager.taxonomy import ClipVariantType
+from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
+from invokeai.backend.util.logging import InvokeAILogger
+
+logger = InvokeAILogger.get_logger()
+
+
+def _fast_safetensors_reader(path: str) -> Dict[str, torch.Tensor]:
+    checkpoint = {}
+    device = torch.device("meta")
+    with open(path, "rb") as f:
+        definition_len = int.from_bytes(f.read(8), "little")
+        definition_json = f.read(definition_len)
+        definition = json.loads(definition_json)
+
+        if "__metadata__" in definition and definition["__metadata__"].get("format", "pt") not in {
+            "pt",
+            "torch",
+            "pytorch",
+        }:
+            raise Exception("Supported only pytorch safetensors files")
+        definition.pop("__metadata__", None)
+
+        for key, info in definition.items():
+            dtype = {
+                "I8": torch.int8,
+                "I16": torch.int16,
+                "I32": torch.int32,
+                "I64": torch.int64,
+                "F16": torch.float16,
+                "F32": torch.float32,
+                "F64": torch.float64,
+            }[info["dtype"]]
+
+            checkpoint[key] = torch.empty(info["shape"], dtype=dtype, device=device)
+
+    return checkpoint
+
+
+def read_checkpoint_meta(path: Union[str, Path], scan: bool = True) -> Dict[str, torch.Tensor]:
+    if str(path).endswith(".safetensors"):
+        try:
+            path_str = path.as_posix() if isinstance(path, Path) else path
+            checkpoint = _fast_safetensors_reader(path_str)
+        except Exception:
+            # TODO: create issue for support "meta"?
+            checkpoint = safetensors.torch.load_file(path, device="cpu")
+    elif str(path).endswith(".gguf"):
+        # The GGUF reader used here uses numpy memmap, so these tensors are not loaded into memory during this function
+        checkpoint = gguf_sd_loader(Path(path), compute_dtype=torch.float32)
+    else:
+        if scan:
+            scan_result = pscan.scan_file_path(path)
+            if scan_result.infected_files != 0:
+                if get_config().unsafe_disable_picklescan:
+                    logger.warning(
+                        f"The model {path} is potentially infected by malware, but picklescan is disabled. "
+                        "Proceeding with caution."
+                    )
+                else:
+                    raise RuntimeError(f"The model {path} is potentially infected by malware. Aborting import.")
+            if scan_result.scan_err:
+                if get_config().unsafe_disable_picklescan:
+                    logger.warning(
+                        f"Error scanning the model at {path} for malware, but picklescan is disabled. "
+                        "Proceeding with caution."
+                    )
+                else:
+                    raise RuntimeError(f"Error scanning the model at {path} for malware. Aborting import.")
+
+        checkpoint = torch.load(path, map_location=torch.device("meta"))
+    return checkpoint
+
+
+def lora_token_vector_length(checkpoint: dict[str | int, torch.Tensor]) -> Optional[int]:
+    """
+    Given a checkpoint in memory, return the lora token vector length
+
+    :param checkpoint: The checkpoint
+    """
+
+    def _get_shape_1(key: str, tensor: torch.Tensor, checkpoint: dict[str | int, torch.Tensor]) -> Optional[int]:
+        lora_token_vector_length = None
+
+        if "." not in key:
+            return lora_token_vector_length  # wrong key format
+        model_key, lora_key = key.split(".", 1)
+
+        # check lora/locon
+        if lora_key == "lora_down.weight":
+            lora_token_vector_length = tensor.shape[1]
+
+        # check loha (don't worry about hada_t1/hada_t2 as it used only in 4d shapes)
+        elif lora_key in ["hada_w1_b", "hada_w2_b"]:
+            lora_token_vector_length = tensor.shape[1]
+
+        # check lokr (don't worry about lokr_t2 as it used only in 4d shapes)
+        elif "lokr_" in lora_key:
+            if model_key + ".lokr_w1" in checkpoint:
+                _lokr_w1 = checkpoint[model_key + ".lokr_w1"]
+            elif model_key + "lokr_w1_b" in checkpoint:
+                _lokr_w1 = checkpoint[model_key + ".lokr_w1_b"]
+            else:
+                return lora_token_vector_length  # unknown format
+
+            if model_key + ".lokr_w2" in checkpoint:
+                _lokr_w2 = checkpoint[model_key + ".lokr_w2"]
+            elif model_key + "lokr_w2_b" in checkpoint:
+                _lokr_w2 = checkpoint[model_key + ".lokr_w2_b"]
+            else:
+                return lora_token_vector_length  # unknown format
+
+            lora_token_vector_length = _lokr_w1.shape[1] * _lokr_w2.shape[1]
+
+        elif lora_key == "diff":
+            lora_token_vector_length = tensor.shape[1]
+
+        # ia3 can be detected only by shape[0] in text encoder
+        elif lora_key == "weight" and "lora_unet_" not in model_key:
+            lora_token_vector_length = tensor.shape[0]
+
+        return lora_token_vector_length
+
+    lora_token_vector_length = None
+    lora_te1_length = None
+    lora_te2_length = None
+    for key, tensor in checkpoint.items():
+        if isinstance(key, int):
+            continue
+        if key.startswith("lora_unet_") and ("_attn2_to_k." in key or "_attn2_to_v." in key):
+            lora_token_vector_length = _get_shape_1(key, tensor, checkpoint)
+        elif key.startswith("lora_unet_") and (
+            "time_emb_proj.lora_down" in key
+        ):  # recognizes format at https://civitai.com/models/224641
+            lora_token_vector_length = _get_shape_1(key, tensor, checkpoint)
+        elif key.startswith("lora_te") and "_self_attn_" in key:
+            tmp_length = _get_shape_1(key, tensor, checkpoint)
+            if key.startswith("lora_te_"):
+                lora_token_vector_length = tmp_length
+            elif key.startswith("lora_te1_"):
+                lora_te1_length = tmp_length
+            elif key.startswith("lora_te2_"):
+                lora_te2_length = tmp_length
+
+        if lora_te1_length is not None and lora_te2_length is not None:
+            lora_token_vector_length = lora_te1_length + lora_te2_length
+
+        if lora_token_vector_length is not None:
+            break
+
+    return lora_token_vector_length
+
+
+def convert_bundle_to_flux_transformer_checkpoint(
+    transformer_state_dict: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    original_state_dict: dict[str, torch.Tensor] = {}
+    keys_to_remove: list[str] = []
+
+    for k, v in transformer_state_dict.items():
+        if not k.startswith("model.diffusion_model"):
+            keys_to_remove.append(k)  # This can be removed in the future if we only want to delete transformer keys
+            continue
+        if k.endswith("scale"):
+            # Scale math must be done at bfloat16 due to our current flux model
+            # support limitations at inference time
+            v = v.to(dtype=torch.bfloat16)
+        new_key = k.replace("model.diffusion_model.", "")
+        original_state_dict[new_key] = v
+        keys_to_remove.append(k)
+
+    # Remove processed keys from the original dictionary, leaving others in case
+    # other model state dicts need to be pulled
+    for k in keys_to_remove:
+        del transformer_state_dict[k]
+
+    return original_state_dict
+
+
+def get_clip_variant_type(location: str) -> Optional[ClipVariantType]:
+    try:
+        path = Path(location)
+        config_path = path / "config.json"
+        if not config_path.exists():
+            config_path = path / "text_encoder" / "config.json"
+        if not config_path.exists():
+            return ClipVariantType.L
+        with open(config_path) as file:
+            clip_conf = json.load(file)
+            hidden_size = clip_conf.get("hidden_size", -1)
+            match hidden_size:
+                case 1280:
+                    return ClipVariantType.G
+                case 768:
+                    return ClipVariantType.L
+                case _:
+                    return ClipVariantType.L
+    except Exception:
+        return ClipVariantType.L

invokeai/backend/model_manager/util/select_hf_files.py

@@ -0,0 +1,218 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+"""
+Select the files from a HuggingFace repository needed for a particular model variant.
+
+Usage:
+```
+from invokeai.backend.model_manager.util.select_hf_files import select_hf_model_files
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch
+
+metadata = HuggingFaceMetadataFetch().from_url("https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0")
+files_to_download = select_hf_model_files(metadata.files, variant='onnx')
+```
+"""
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Set
+
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+
+
+def filter_files(
+    files: List[Path],
+    variant: Optional[ModelRepoVariant] = None,
+    subfolder: Optional[Path] = None,
+    subfolders: Optional[List[Path]] = None,
+) -> List[Path]:
+    """
+    Take a list of files in a HuggingFace repo root and return paths to files needed to load the model.
+
+    :param files: List of files relative to the repo root.
+    :param subfolder: Filter by the indicated subfolder (deprecated, use subfolders instead).
+    :param subfolders: Filter by multiple subfolders. Files from any of these subfolders will be included.
+    :param variant: Filter by files belonging to a particular variant, such as fp16.
+
+    The file list can be obtained from the `files` field of HuggingFaceMetadata,
+    as defined in `invokeai.backend.model_manager.metadata.metadata_base`.
+    """
+    variant = variant or ModelRepoVariant.Default
+    paths: List[Path] = []
+
+    if not files:
+        return []
+
+    root = files[0].parts[0] if files[0].parts else Path(".")
+
+    # Build list of subfolders to filter by
+    filter_subfolders: List[Path] = []
+    if subfolders:
+        filter_subfolders = subfolders
+    elif subfolder:
+        filter_subfolders = [subfolder]
+
+    # if the subfolder is a single file, then bypass the selection and just return it
+    if len(filter_subfolders) == 1:
+        sf = filter_subfolders[0]
+        if sf.suffix in [".safetensors", ".bin", ".onnx", ".xml", ".pth", ".pt", ".ckpt", ".msgpack"]:
+            return [root / sf]
+
+    # Start by filtering on model file extensions, discarding images, docs, etc
+    for file in files:
+        if file.name.endswith((".json", ".txt", ".jinja")):  # .jinja for chat templates
+            paths.append(file)
+        elif file.name.endswith(
+            (
+                "learned_embeds.bin",
+                "ip_adapter.bin",
+                "lora_weights.safetensors",
+                "weights.pb",
+                "onnx_data",
+                "spiece.model",  # Added for `black-forest-labs/FLUX.1-schnell`.
+            )
+        ):
+            paths.append(file)
+        # BRITTLENESS WARNING!!
+        # Diffusers models always seem to have "model" in their name, and the regex filter below is applied to avoid
+        # downloading random checkpoints that might also be in the repo. However there is no guarantee
+        # that a checkpoint doesn't contain "model" in its name, and no guarantee that future diffusers models
+        # will adhere to this naming convention, so this is an area to be careful of.
+        elif re.search(r"model.*\.(safetensors|bin|onnx|xml|pth|pt|ckpt|msgpack)$", file.name):
+            paths.append(file)
+
+    # limit search to subfolder(s) if requested
+    if filter_subfolders:
+        absolute_subfolders = [root / sf for sf in filter_subfolders]
+        paths = [x for x in paths if any(Path(sf) in x.parents for sf in absolute_subfolders)]
+
+    # _filter_by_variant uniquifies the paths and returns a set
+    return sorted(_filter_by_variant(paths, variant))
+
+
+@dataclass
+class SubfolderCandidate:
+    path: Path
+    score: int
+
+
+def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path]:
+    """Select the proper variant files from a list of HuggingFace repo_id paths."""
+    result: set[Path] = set()
+    subfolder_weights: dict[Path, list[SubfolderCandidate]] = {}
+    safetensors_detected = False
+    for path in files:
+        if path.suffix in [".onnx", ".pb", ".onnx_data"]:
+            if variant == ModelRepoVariant.ONNX:
+                result.add(path)
+
+        elif "openvino_model" in path.name:
+            if variant == ModelRepoVariant.OpenVINO:
+                result.add(path)
+
+        elif "flax_model" in path.name:
+            if variant == ModelRepoVariant.Flax:
+                result.add(path)
+
+        # Note: '.model' was added to support:
+        # https://huggingface.co/black-forest-labs/FLUX.1-schnell/blob/768d12a373ed5cc9ef9a9dea7504dc09fcc14842/tokenizer_2/spiece.model
+        # Note: '.jinja' was added to support chat templates for FLUX.2 Klein models
+        elif path.suffix in [".json", ".txt", ".model", ".jinja"]:
+            result.add(path)
+
+        elif variant in [
+            ModelRepoVariant.FP16,
+            ModelRepoVariant.FP32,
+            ModelRepoVariant.Default,
+        ] and path.suffix in [".bin", ".safetensors", ".pt", ".ckpt"]:
+            # For weights files, we want to select the best one for each subfolder. For example, we may have multiple
+            # text encoders:
+            #
+            # - text_encoder/model.fp16.safetensors
+            # - text_encoder/model.safetensors
+            # - text_encoder/pytorch_model.bin
+            # - text_encoder/pytorch_model.fp16.bin
+            #
+            # We prefer safetensors over other file formats and an exact variant match. We'll score each file based on
+            # variant and format and select the best one.
+
+            if safetensors_detected and path.suffix == ".bin":
+                continue
+
+            parent = path.parent
+            score = 0
+
+            if path.suffix == ".safetensors":
+                safetensors_detected = True
+                if parent in subfolder_weights:
+                    subfolder_weights[parent] = [sfc for sfc in subfolder_weights[parent] if sfc.path.suffix != ".bin"]
+                score += 1
+
+            candidate_variant_label = path.suffixes[0] if len(path.suffixes) == 2 else None
+
+            # Some special handling is needed here if there is not an exact match and if we cannot infer the variant
+            # from the file name. In this case, we only give this file a point if the requested variant is FP32 or DEFAULT.
+            if (
+                variant is not ModelRepoVariant.Default
+                and candidate_variant_label
+                and candidate_variant_label.startswith(f".{variant.value}")
+            ) or (not candidate_variant_label and variant in [ModelRepoVariant.FP32, ModelRepoVariant.Default]):
+                score += 1
+
+            if parent not in subfolder_weights:
+                subfolder_weights[parent] = []
+
+            subfolder_weights[parent].append(SubfolderCandidate(path=path, score=score))
+
+        else:
+            continue
+
+    for candidate_list in subfolder_weights.values():
+        # Check if at least one of the files has the explicit fp16 variant.
+        at_least_one_fp16 = False
+        for candidate in candidate_list:
+            if len(candidate.path.suffixes) == 2 and candidate.path.suffixes[0].startswith(".fp16"):
+                at_least_one_fp16 = True
+                break
+
+        if not at_least_one_fp16:
+            # If none of the candidates in this candidate_list have the explicit fp16 variant label, then this
+            # candidate_list probably doesn't adhere to the variant naming convention that we expected. In this case,
+            # we'll simply keep all the candidates. An example of a model that hits this case is
+            # `black-forest-labs/FLUX.1-schnell` (as of commit 012d2fd).
+            for candidate in candidate_list:
+                result.add(candidate.path)
+
+        # The candidate_list seems to have the expected variant naming convention. We'll select the highest scoring
+        # candidate.
+        highest_score_candidate = max(candidate_list, key=lambda candidate: candidate.score)
+        if highest_score_candidate:
+            pattern = r"^(.*?)-\d+-of-\d+(\.\w+)$"
+            match = re.match(pattern, highest_score_candidate.path.as_posix())
+            if match:
+                for candidate in candidate_list:
+                    if candidate.path.as_posix().startswith(match.group(1)) and candidate.path.as_posix().endswith(
+                        match.group(2)
+                    ):
+                        result.add(candidate.path)
+            else:
+                result.add(highest_score_candidate.path)
+
+    # If one of the architecture-related variants was specified and no files matched other than
+    # config and text files then we return an empty list
+    if (
+        variant
+        and variant in [ModelRepoVariant.ONNX, ModelRepoVariant.OpenVINO, ModelRepoVariant.Flax]
+        and not any(variant.value in x.name for x in result)
+    ):
+        return set()
+
+    # Prune folders that contain just a `config.json`. This happens when
+    # the requested variant (e.g. "onnx") is missing
+    directories: Dict[Path, int] = {}
+    for x in result:
+        if not x.parent:
+            continue
+        directories[x.parent] = directories.get(x.parent, 0) + 1
+
+    return {x for x in result if directories[x.parent] > 1 or x.name != "config.json"}