Download utils/models.py from turtle89431/Moxie-Multimedia: direct link, hf CLI and curl.
- Browser
- Download file 28.5 kB
-
https://huggingface.co/turtle89431/Moxie-Multimedia/resolve/main/utils/models.py
- Command line
-
hf download hf://turtle89431/Moxie-Multimedia/utils/models.py
-
curl -L -o models.py https://huggingface.co/turtle89431/Moxie-Multimedia/resolve/main/utils/models.py
28.5 kB
| 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] = {} | |
| class TurboModelDetection: | |
| status: Literal["turbo", "non_turbo", "unknown"] | |
| source: Literal[ | |
| "model_patches", "model_metadata", "graph_prompt", "fallback" | |
| ] | |
| evidence: str | |
| patch_count: int = 0 | |
| 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 | |
| class EasyMediaModel: | |
| name: str | |
| display_name: str | |
| category: str | |
| filename: str | |
| url: str | |
| urls: tuple[str, ...] = field(default_factory=tuple) | |
| def directory(self) -> Path: | |
| return Path(folder_paths.models_dir) / self.category | |
| 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 | |