utils/models.py

from __future__ import annotations

import asyncio
import uuid
from collections import deque
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal

import aiohttp
import folder_paths


MODEL_MISSING_EVENT = "easy-media-model-missing"
MODEL_DOWNLOAD_TIMEOUT_SECONDS = 600
_MODEL_DOWNLOAD_LOCKS: dict[str, asyncio.Lock] = {}


@dataclass(frozen=True)
class TurboModelDetection:
    status: Literal["turbo", "non_turbo", "unknown"]
    source: Literal[
        "model_patches", "model_metadata", "graph_prompt", "fallback"
    ]
    evidence: str
    patch_count: int = 0

    @property
    def is_turbo(self) -> bool:
        return self.status == "turbo"

    def as_dict(self) -> dict[str, str | int | bool]:
        return {
            "status": self.status,
            "is_turbo": self.is_turbo,
            "source": self.source,
            "evidence": self.evidence,
            "patch_count": self.patch_count,
        }


def _iter_model_patchers(model: Any):
    if isinstance(model, (list, tuple)):
        for item in model:
            yield from _iter_model_patchers(item)
        return
    if model is not None:
        yield model


def _lora_patch_rank(patch_value: Any) -> int | None:
    if not isinstance(patch_value, (list, tuple)) or len(patch_value) < 2:
        return None
    adapter = patch_value[1]
    weights = getattr(adapter, "weights", None)
    if not isinstance(weights, (list, tuple)) or len(weights) < 2:
        return None
    down_weight = weights[1]
    shape = getattr(down_weight, "shape", None)
    if shape is None or len(shape) < 1:
        return None
    try:
        return int(shape[0])
    except (TypeError, ValueError):
        return None


def _patch_has_type(patch_values: Any, patch_type: str) -> bool:
    if not isinstance(patch_values, (list, tuple)):
        return False
    for patch_value in patch_values:
        if not isinstance(patch_value, (list, tuple)) or len(patch_value) < 2:
            continue
        adapter = patch_value[1]
        if (
            isinstance(adapter, (list, tuple))
            and len(adapter) > 0
            and adapter[0] == patch_type
        ):
            return True
    return False


def _model_patch_evidence(
    model: Any,
) -> tuple[int, list[str], set[int]]:
    patch_count = 0
    patch_keys: list[str] = []
    patch_ranks: set[int] = set()
    for patcher in _iter_model_patchers(model):
        patches = getattr(patcher, "patches", None)
        if not isinstance(patches, Mapping):
            continue
        for key, patch_values in patches.items():
            if isinstance(patch_values, (list, tuple)):
                count = len(patch_values)
            else:
                count = 0 if patch_values is None else 1
            if count <= 0:
                continue
            patch_count += count
            if isinstance(patch_values, (list, tuple)):
                for patch_value in patch_values:
                    rank = _lora_patch_rank(patch_value)
                    if rank is not None:
                        patch_ranks.add(rank)
            if len(patch_keys) < 5:
                patch_keys.append(str(key))
    return patch_count, patch_keys, patch_ranks


def _find_turbo_keyword(
    value: Any,
    path: str,
    *,
    depth: int = 0,
    visited: set[int] | None = None,
) -> tuple[str, str] | None:
    if depth > 6:
        return None
    if visited is None:
        visited = set()

    if isinstance(value, bytes):
        text = value.decode("utf-8", errors="replace")
        return (path, text) if "turbo" in text.lower() else None
    if isinstance(value, (str, Path)):
        text = str(value)
        return (path, text) if "turbo" in text.lower() else None

    if isinstance(value, Mapping):
        value_id = id(value)
        if value_id in visited:
            return None
        visited.add(value_id)
        for key, item in value.items():
            key_text = str(key)
            key_path = f"{path}.{key_text}"
            if "turbo" in key_text.lower():
                return key_path, key_text
            match = _find_turbo_keyword(
                item,
                key_path,
                depth=depth + 1,
                visited=visited,
            )
            if match is not None:
                return match
        return None

    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
        value_id = id(value)
        if value_id in visited:
            return None
        visited.add(value_id)
        for index, item in enumerate(value):
            match = _find_turbo_keyword(
                item,
                f"{path}[{index}]",
                depth=depth + 1,
                visited=visited,
            )
            if match is not None:
                return match
    return None


def _model_metadata_sources(model: Any):
    for index, patcher in enumerate(_iter_model_patchers(model)):
        prefix = f"model[{index}]"
        yield f"{prefix}.attachments", getattr(patcher, "attachments", None)
        yield f"{prefix}.model_options", getattr(patcher, "model_options", None)

        base_model = getattr(patcher, "model", None)
        if base_model is None:
            continue
        yield f"{prefix}.model.metadata", getattr(base_model, "metadata", None)
        yield (
            f"{prefix}.model.model_metadata",
            getattr(base_model, "model_metadata", None),
        )
        model_config = getattr(base_model, "model_config", None)
        if model_config is not None:
            yield (
                f"{prefix}.model.model_config",
                getattr(model_config, "__dict__", None),
            )


def _find_sampler_steps(
    value: Any,
    path: str,
    *,
    depth: int = 0,
    visited: set[int] | None = None,
) -> tuple[str, int] | None:
    if depth > 6 or not isinstance(value, Mapping):
        return None
    if visited is None:
        visited = set()
    value_id = id(value)
    if value_id in visited:
        return None
    visited.add(value_id)

    for key, item in value.items():
        key_text = str(key)
        key_path = f"{path}.{key_text}"
        if key_text.lower() == "sampler_steps":
            try:
                return key_path, int(item)
            except (TypeError, ValueError):
                continue
        match = _find_sampler_steps(
            item,
            key_path,
            depth=depth + 1,
            visited=visited,
        )
        if match is not None:
            return match
    return None


def _find_pdd_effective_steps(
    value: Any,
    path: str,
    *,
    depth: int = 0,
    visited: set[int] | None = None,
) -> tuple[str, int, int, int] | None:
    if depth > 6 or not isinstance(value, Mapping):
        return None
    if visited is None:
        visited = set()
    value_id = id(value)
    if value_id in visited:
        return None
    visited.add(value_id)

    normalized = {str(key).lower(): item for key, item in value.items()}
    if "pdd_num_steps" in normalized and "pdd_block_size" in normalized:
        try:
            num_steps = int(normalized["pdd_num_steps"])
            block_size = int(normalized["pdd_block_size"])
        except (TypeError, ValueError):
            pass
        else:
            if num_steps > 0 and block_size > 0 and num_steps % block_size == 0:
                return path, num_steps, block_size, num_steps // block_size

    for key, item in value.items():
        match = _find_pdd_effective_steps(
            item,
            f"{path}.{key}",
            depth=depth + 1,
            visited=visited,
        )
        if match is not None:
            return match
    return None


def detect_turbo_model(model: Any) -> TurboModelDetection:
    """Detect Turbo using only data retained on the ComfyUI model object."""
    patch_count, patch_keys, patch_ranks = _model_patch_evidence(model)
    if patch_count > 0:
        patcher_entries: dict[str, list[Any]] = {}
        for patcher in _iter_model_patchers(model):
            patches = getattr(patcher, "patches", None)
            if not isinstance(patches, Mapping):
                continue
            for key, values in patches.items():
                if values is None:
                    continue
                normalized_values = (
                    list(values)
                    if isinstance(values, (list, tuple))
                    else [values]
                )
                patcher_entries.setdefault(str(key).lower(), []).extend(
                    normalized_values
                )
        patcher_keys = list(patcher_entries)
        has_attention = any(
            ".attn.qkv_proj." in key or ".attn.out_proj." in key
            for key in patcher_keys
        )
        has_adaln = any(".adaln_proj." in key for key in patcher_keys)
        has_mlp_fc1 = any(".mlp.fc1." in key for key in patcher_keys)
        has_mlp_fc2 = any(".mlp.fc2." in key for key in patcher_keys)
        pdd_output_targets = (
            "diffusion_model.final_layer.audio_out.weight",
            "diffusion_model.final_layer.video_out.weight",
        )
        has_pdd_output_heads = all(
            _patch_has_type(patcher_entries.get(key), "set")
            for key in pdd_output_targets
        )

        has_shared_turbo_targets = has_attention and has_mlp_fc1 and has_mlp_fc2
        matches_pdd_eight_step = (
            has_shared_turbo_targets and has_adaln and has_pdd_output_heads
        )
        matches_four_step = (
            has_shared_turbo_targets and has_adaln and 64 in patch_ranks
        )
        matches_lightx2v_eight_step = (
            has_shared_turbo_targets
            and not has_adaln
            and {128, 384}.issubset(patch_ranks)
        )
        if (
            matches_pdd_eight_step
            or matches_four_step
            or matches_lightx2v_eight_step
        ):
            if matches_pdd_eight_step:
                fingerprint_name = "PDD 8-step output-head/attention/MLP/AdaLN"
            elif matches_four_step:
                fingerprint_name = "4-step attention/MLP/AdaLN"
            else:
                fingerprint_name = "LightX2V 8-step attention/MLP"
            return TurboModelDetection(
                status="turbo",
                source="model_patches",
                evidence=(
                    f"Moxie Turbo {fingerprint_name} fingerprint matched with "
                    f"ranks {sorted(patch_ranks)}"
                ),
                patch_count=patch_count,
            )

        if (
            has_attention
            and not has_adaln
            and not has_mlp_fc1
            and not has_mlp_fc2
            and patch_ranks == {32}
        ):
            return TurboModelDetection(
                status="non_turbo",
                source="model_patches",
                evidence=(
                    "attention-only rank-32 LoRA fingerprint matched; this is not "
                    "the Moxie Turbo 4-step structure"
                ),
                patch_count=patch_count,
            )

    for path, metadata in _model_metadata_sources(model):
        pdd_steps = _find_pdd_effective_steps(metadata, path)
        if pdd_steps is not None:
            pdd_path, num_steps, block_size, effective_steps = pdd_steps
            if effective_steps <= 8:
                return TurboModelDetection(
                    status="turbo",
                    source="model_metadata",
                    evidence=(
                        f"PDD Turbo metadata found at {pdd_path}: "
                        f"pdd_num_steps={num_steps}, pdd_block_size={block_size}, "
                        f"effective_steps={effective_steps}"
                    ),
                    patch_count=patch_count,
                )
        sampler_steps = _find_sampler_steps(metadata, path)
        if sampler_steps is not None:
            steps_path, steps = sampler_steps
            if steps <= 8:
                return TurboModelDetection(
                    status="turbo",
                    source="model_metadata",
                    evidence=f"{steps_path}={steps}",
                    patch_count=patch_count,
                )
        match = _find_turbo_keyword(metadata, path)
        if match is None:
            continue
        match_path, match_value = match
        return TurboModelDetection(
            status="turbo",
            source="model_metadata",
            evidence=f"turbo keyword found at {match_path}: {match_value[:160]}",
            patch_count=patch_count,
        )

    patch_summary = ""
    if patch_count > 0:
        key_summary = ", ".join(patch_keys) if patch_keys else "unavailable"
        patch_summary = (
            f"; {patch_count} unclassified patch entries with ranks "
            f"{sorted(patch_ranks)}, sample keys: {key_summary}"
        )
    return TurboModelDetection(
        status="unknown",
        source="fallback",
        evidence=(
            "no Turbo patch fingerprint, sampler_steps<=8, or turbo keyword in "
            f"model-side metadata{patch_summary}"
        ),
        patch_count=patch_count,
    )


def _prompt_mapping(prompt: Any) -> Mapping[Any, Any] | None:
    while isinstance(prompt, (list, tuple)) and len(prompt) == 1:
        prompt = prompt[0]
    return prompt if isinstance(prompt, Mapping) else None


def _prompt_node(
    prompt: Mapping[Any, Any], node_id: Any
) -> tuple[str, Mapping[Any, Any]] | None:
    normalized_id = str(node_id)
    node = prompt.get(normalized_id)
    if node is None:
        node = prompt.get(node_id)
    if not isinstance(node, Mapping):
        return None
    return normalized_id, node


def _linked_node_id(value: Any, prompt: Mapping[Any, Any]) -> str | None:
    if not isinstance(value, (list, tuple)) or len(value) != 2:
        return None
    if not isinstance(value[1], int):
        return None
    linked_node = _prompt_node(prompt, value[0])
    return linked_node[0] if linked_node is not None else None


def _input_links(value: Any, prompt: Mapping[Any, Any]):
    linked_node_id = _linked_node_id(value, prompt)
    if linked_node_id is not None:
        yield linked_node_id
        return
    if isinstance(value, Mapping):
        for item in value.values():
            yield from _input_links(item, prompt)
    elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
        for item in value:
            yield from _input_links(item, prompt)


def _upstream_nodes(
    prompt: Mapping[Any, Any], start_node_ids: Sequence[str]
):
    queue = deque(start_node_ids)
    visited: set[str] = set()
    while queue:
        node_id = queue.popleft()
        if node_id in visited:
            continue
        visited.add(node_id)
        prompt_node = _prompt_node(prompt, node_id)
        if prompt_node is None:
            continue
        normalized_id, node = prompt_node
        yield normalized_id, node
        inputs = node.get("inputs")
        if not isinstance(inputs, Mapping):
            continue
        for linked_node_id in _input_links(inputs, prompt):
            if linked_node_id not in visited:
                queue.append(linked_node_id)


TURBO_KEYWORDS = ("turbo", "acc", "8step", "4step", "3step")

def _has_turbo_name(value: Any) -> bool:
    if not isinstance(value, (str, Path)):
        return False
    val_lower = str(value).lower()
    return any(kw in val_lower for kw in TURBO_KEYWORDS)


def _is_enabled(value: Any) -> bool:
    if value is True:
        return True
    return isinstance(value, str) and value.strip().lower() == "true"


def detect_turbo_lora_from_prompt(
    prompt: Any, current_node_id: Any
) -> TurboModelDetection | None:
    """Find a Turbo LoRA filename upstream when the model retained no metadata."""
    prompt_graph = _prompt_mapping(prompt)
    if prompt_graph is None:
        return None
    current_node = _prompt_node(prompt_graph, current_node_id)
    if current_node is None:
        return None
    current_inputs = current_node[1].get("inputs")
    if not isinstance(current_inputs, Mapping):
        return None
    loader_node_id = _linked_node_id(
        current_inputs.get("model_loader"), prompt_graph
    )
    if loader_node_id is None:
        return None
    loader_node = _prompt_node(prompt_graph, loader_node_id)
    if loader_node is None:
        return None

    loader_class = loader_node[1].get("class_type")
    if loader_class == "easy modelLoaderPack":
        loader_inputs = loader_node[1].get("inputs")
        if not isinstance(loader_inputs, Mapping):
            return None
        model_node_id = _linked_node_id(loader_inputs.get("model"), prompt_graph)
        if model_node_id is None:
            return None
        for node_id, node in _upstream_nodes(prompt_graph, [model_node_id]):
            if node.get("class_type") != "LoraLoaderModelOnly":
                continue
            inputs = node.get("inputs")
            if not isinstance(inputs, Mapping):
                continue
            lora_name = inputs.get("lora_name")
            if _has_turbo_name(lora_name):
                return TurboModelDetection(
                    status="turbo",
                    source="graph_prompt",
                    evidence=(
                        f"upstream LoraLoaderModelOnly node {node_id} uses "
                        f"Turbo LoRA: {lora_name}"
                    ),
                )
        return None

    for node_id, node in _upstream_nodes(prompt_graph, [loader_node_id]):
        class_type = node.get("class_type")

        # fast lorasLoader — check lora_1..lora_10 entries
        if class_type == "fast lorasLoader":
            inputs = node.get("inputs")
            if not isinstance(inputs, Mapping):
                continue
            for index in range(1, 11):
                key = f"lora_{index}"
                if key not in inputs:
                    break
                lora_config = inputs[key]
                if not isinstance(lora_config, Mapping):
                    continue
                lora_name = lora_config.get("lora")
                if _has_turbo_name(lora_name) and _is_enabled(
                    lora_config.get("enabled")
                ):
                    return TurboModelDetection(
                        status="turbo",
                        source="graph_prompt",
                        evidence=(
                            f"upstream fast lorasLoader node {node_id} input {key} "
                            f"enables Turbo LoRA: {lora_name}"
                        ),
                    )
            continue

        # UNETLoader — check unet_name for turbo
        if class_type == "UNETLoader":
            inputs = node.get("inputs")
            if isinstance(inputs, Mapping):
                unet_name = inputs.get("unet_name")
                if _has_turbo_name(unet_name):
                    return TurboModelDetection(
                        status="turbo",
                        source="graph_prompt",
                        evidence=(
                            f"upstream UNETLoader node {node_id} uses "
                            f"Turbo UNET: {unet_name}"
                        ),
                    )
            continue

        # fast h3Loader — check 主模型 and 副模型 for turbo
        if class_type == "fast h3Loader":
            inputs = node.get("inputs")
            if not isinstance(inputs, Mapping):
                continue
            for model_key in ("主模型", "副模型"):
                model_value = inputs.get(model_key)
                if _has_turbo_name(model_value):
                    return TurboModelDetection(
                        status="turbo",
                        source="graph_prompt",
                        evidence=(
                            f"upstream fast h3Loader node {node_id} "
                            f"{model_key} is Turbo: {model_value}"
                        ),
                    )
    return None


@dataclass(frozen=True)
class EasyMediaModel:
    name: str
    display_name: str
    category: str
    filename: str
    url: str
    urls: tuple[str, ...] = field(default_factory=tuple)

    @property
    def directory(self) -> Path:
        return Path(folder_paths.models_dir) / self.category

    @property
    def path(self) -> Path:
        return self.directory / self.filename


MODEL_REGISTRY: dict[str, EasyMediaModel] = {
    "omnishotcut": EasyMediaModel(
        name="omnishotcut",
        display_name="OmniShotCut",
        category="checkpoints",
        filename="OmniShotCut_ckpt.pth",
        url="https://huggingface.co/uva-cv-lab/OmniShotCut/resolve/main/OmniShotCut_ckpt.pth",
    ),
    "qwen3-asr": EasyMediaModel(
        name="qwen3-asr",
        display_name="Qwen3-ASR",
        category="",
        filename="Qwen3-ASR",
        url="https://huggingface.co/Qwen/Qwen3-ASR-1.7B",
        urls=(
            "https://huggingface.co/Qwen/Qwen3-ASR-1.7B",
            "https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B",
        ),
    ),
    "whisper-large-v3": EasyMediaModel(
        name="whisper-large-v3",
        display_name="Whisper Large V3",
        category="audio_encoders",
        filename="whisper_large_v3_fp16.safetensors",
        url="https://huggingface.co/Comfy-Org/HuMo_ComfyUI/resolve/main/split_files/audio_encoders/whisper_large_v3_fp16.safetensors",
    ),
    "voxcpm2": EasyMediaModel(
        name="voxcpm2",
        display_name="VoxCPM2",
        category="voxcpm",
        filename="VoxCPM2",
        url="https://huggingface.co/openbmb/VoxCPM2",
    ),
}


class MissingEasyMediaModelError(FileNotFoundError):
    def __init__(self, model: EasyMediaModel):
        self.model = model
        super().__init__(
            f"{model.display_name} model is not installed. "
            f"Download {model.filename} to {model.directory}."
        )


def get_model_info(model_name: str) -> EasyMediaModel:
    try:
        return MODEL_REGISTRY[model_name]
    except KeyError as error:
        raise ValueError(f"Unknown Easy Media model: {model_name}") from error


def get_model_path(model_name: str) -> Path:
    """Return the expected local path for a registered Easy Media model."""
    return get_model_info(model_name).path


def model_payload(model: EasyMediaModel) -> dict:
    payload = {
        "name": model.name,
        "display_name": model.display_name,
        "filename": model.filename,
        "directory": str(model.directory),
        "path": str(model.path),
        "url": model.url,
    }
    if model.urls:
        payload["urls"] = list(model.urls)
    return payload


def notify_missing_model(model_name: str) -> dict:
    model = get_model_info(model_name)
    payload = model_payload(model)
    try:
        from server import PromptServer

        PromptServer.instance.send_sync(MODEL_MISSING_EVENT, payload)
    except Exception as error:
        print(f"[Moxie] Failed to notify missing model {model_name}: {error}")
    return payload


def require_model_path(model_name: str) -> Path:
    model = get_model_info(model_name)
    if model.path.is_file():
        return model.path
    notify_missing_model(model_name)
    raise MissingEasyMediaModelError(model)


async def download_model(model_name: str) -> Path:
    model = get_model_info(model_name)
    target = model.path
    lock = _MODEL_DOWNLOAD_LOCKS.setdefault(model.name, asyncio.Lock())

    async with lock:
        if model.name == "qwen3-asr":
            return await _download_qwen3_asr_bundle(model)
        if model.name == "voxcpm2":
            return await _download_snapshot_model(model, "openbmb/VoxCPM2")

        if target.is_file():
            return target

        model.directory.mkdir(parents=True, exist_ok=True)
        partial = target.with_name(f"{target.name}.{uuid.uuid4().hex}.download")
        timeout = aiohttp.ClientTimeout(total=MODEL_DOWNLOAD_TIMEOUT_SECONDS)

        try:
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.get(model.url) as response:
                    response.raise_for_status()
                    with partial.open("wb") as file:
                        while True:
                            chunk = await response.content.read(1024 * 1024)
                            if not chunk:
                                break
                            file.write(chunk)
            partial.replace(target)
            return target
        except asyncio.TimeoutError as error:
            partial.unlink(missing_ok=True)
            raise TimeoutError(
                f"Automatic download timed out after {MODEL_DOWNLOAD_TIMEOUT_SECONDS} seconds."
            ) from error
        except Exception:
            partial.unlink(missing_ok=True)
            raise
        raise


def require_qwen_asr_model_dirs() -> tuple[Path, Path]:
    """Return the ASR and aligner model directories, raising if either is missing."""
    root = Path(folder_paths.models_dir) / "Qwen3-ASR"
    candidates = ("Qwen3-ASR-1.7B", "Qwen3-ASR-0.6B")
    asr_dir = next((root / name for name in candidates if (root / name).is_dir()), None)
    aligner_dir = root / "Qwen3-ForcedAligner-0.6B"
    if asr_dir is not None and aligner_dir.is_dir():
        return asr_dir, aligner_dir
    notify_missing_model("qwen3-asr")
    raise MissingEasyMediaModelError(get_model_info("qwen3-asr"))


def require_whisper_large_v3_model_path() -> Path:
    """Return the local Whisper Large V3 audio encoder safetensors file."""
    model = get_model_info("whisper-large-v3")
    target_name = "whisper_large_v3"
    preferred_filenames = {
        model.filename.lower(),
        "whisper-large-v3.safetensors",
    }
    candidates = []
    for filename in folder_paths.get_filename_list("audio_encoders"):
        path = Path(filename)
        if path.suffix.lower() != ".safetensors":
            continue
        normalized_filename = filename.lower()
        is_preferred = path.name.lower() in preferred_filenames
        if "encode" in normalized_filename:
            continue
        if not is_preferred and target_name not in normalized_filename:
            continue
        candidates.append(filename)

    candidates.sort(
        key=lambda filename: Path(filename).name.lower() not in preferred_filenames
    )
    for filename in candidates:
        full_path = folder_paths.get_full_path("audio_encoders", filename)
        if full_path:
            return Path(full_path)
    notify_missing_model("whisper-large-v3")
    raise MissingEasyMediaModelError(model)


async def _download_qwen3_asr_bundle(model: EasyMediaModel) -> Path:
    target = model.path
    asr_dir = target / "Qwen3-ASR-1.7B"
    aligner_dir = target / "Qwen3-ForcedAligner-0.6B"
    if asr_dir.is_dir() and aligner_dir.is_dir():
        return target

    try:
        from huggingface_hub import snapshot_download  # type: ignore[import]
    except ImportError as error:
        raise RuntimeError(
            "Automatic Qwen3-ASR download requires huggingface_hub. "
            "Install it with: pip install huggingface_hub"
        ) from error

    target.mkdir(parents=True, exist_ok=True)

    async def download_snapshot(repo_id: str, local_dir: Path) -> None:
        await asyncio.to_thread(
            snapshot_download,
            repo_id=repo_id,
            local_dir=str(local_dir),
            local_dir_use_symlinks=False,
        )

    await download_snapshot("Qwen/Qwen3-ASR-1.7B", asr_dir)
    await download_snapshot("Qwen/Qwen3-ForcedAligner-0.6B", aligner_dir)
    return target


async def _download_snapshot_model(model: EasyMediaModel, repo_id: str) -> Path:
    target = model.path
    if target.is_dir():
        return target

    try:
        from huggingface_hub import snapshot_download  # type: ignore[import]
    except ImportError as error:
        raise RuntimeError(
            f"Automatic {model.display_name} download requires huggingface_hub. "
            "Install it with: pip install huggingface_hub"
        ) from error

    target.mkdir(parents=True, exist_ok=True)
    await asyncio.to_thread(
        snapshot_download,
        repo_id=repo_id,
        local_dir=str(target),
        local_dir_use_symlinks=False,
    )
    return target