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