# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Load Kimodo diffusion models from local checkpoints or Hugging Face.""" from pathlib import Path from typing import Optional from huggingface_hub import snapshot_download from omegaconf import OmegaConf from .loading import ( AVAILABLE_MODELS, DEFAULT_MODEL, DEFAULT_TEXT_ENCODER_URL, MODEL_NAMES, TMR_MODELS, get_env_var, instantiate_from_dict, ) from .registry import get_model_info, resolve_model_name DEFAULT_TEXT_ENCODER = "llm2vec" TEXT_ENCODER_PRESETS = { "llm2vec": { "target": "kimodo.model.LLM2VecEncoder", "kwargs": { "base_model_name_or_path": "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp", "peft_model_name_or_path": "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-supervised", "dtype": "bfloat16", "llm_dim": 4096, "device": "auto", }, } } def _resolve_hf_model_path(modelname: str) -> Path: """Resolve model name to a local path, using Hugging Face cache or CHECKPOINT_DIR.""" try: repo_id = MODEL_NAMES[modelname] except KeyError: raise ValueError(f"Model '{modelname}' not found. Available models: {MODEL_NAMES.keys()}") local_cache = get_env_var("LOCAL_CACHE", "False").lower() == "true" if not local_cache: snapshot_dir = snapshot_download(repo_id=repo_id) # will check online no matter what return Path(snapshot_dir) try: snapshot_dir = snapshot_download(repo_id=repo_id, local_files_only=True) # will check local cache only return Path(snapshot_dir) except Exception: # if local cache is not found, download from online try: snapshot_dir = snapshot_download(repo_id=repo_id) return Path(snapshot_dir) except Exception: raise RuntimeError(f"Could not resolve model '{modelname}' from Hugging Face (repo: {repo_id}). ") from None def _build_api_text_encoder_conf(text_encoder_url: str) -> dict: return { "_target_": "kimodo.model.text_encoder_api.TextEncoderAPI", "url": text_encoder_url, } def _build_local_text_encoder_conf(text_encoder_fp32: bool = False) -> dict: text_encoder_name = get_env_var("TEXT_ENCODER", DEFAULT_TEXT_ENCODER) if text_encoder_name not in TEXT_ENCODER_PRESETS: available = ", ".join(sorted(TEXT_ENCODER_PRESETS)) raise ValueError(f"Unknown TEXT_ENCODER='{text_encoder_name}'. Available: {available}") preset = TEXT_ENCODER_PRESETS[text_encoder_name] if text_encoder_fp32: preset["kwargs"]["dtype"] = "float32" return { "_target_": preset["target"], **preset["kwargs"], } def _select_text_encoder_conf(text_encoder_url: str, text_encoder_fp32: bool = False) -> dict: # TEXT_ENCODER_MODE options: # - "api": force TextEncoderAPI # - "local": force local LLM2VecEncoder # - "auto": try API first, fallback to local if unreachable mode = get_env_var("TEXT_ENCODER_MODE", "auto").lower() if mode == "local": return _build_local_text_encoder_conf(text_encoder_fp32) if mode == "api": return _build_api_text_encoder_conf(text_encoder_url) api_conf = _build_api_text_encoder_conf(text_encoder_url) try: text_encoder = instantiate_from_dict(api_conf) # Probe availability early so inference doesn't fail later. text_encoder(["healthcheck"]) return api_conf except Exception as error: print( "Text encoder service is unreachable, falling back to local LLM2Vec " f"encoder. ({type(error).__name__}: {error})" ) return _build_local_text_encoder_conf(text_encoder_fp32) def load_model( modelname=None, device=None, eval_mode: bool = True, default_family: Optional[str] = "Kimodo", return_resolved_name: bool = False, text_encoder=None, text_encoder_fp32: bool = False, ): """Load a kimodo model by name (e.g. 'g1', 'soma'). Resolution of partial/full names (e.g. Kimodo-SOMA-RP-v1, SOMA) is done inside this function using default_family when the name is not a known short key. Args: modelname: Model identifier; uses DEFAULT_MODEL if None. Can be a short key, a full name (e.g. Kimodo-SOMA-RP-v1), or a partial name; unknown names are resolved via resolve_model_name using default_family. device: Target device for the model (e.g. 'cuda', 'cpu'). eval_mode: If True, set model to eval mode. default_family: Used when modelname is not in AVAILABLE_MODELS to resolve partial names ("Kimodo" for demo/generation, "TMR" for embed script). Default "Kimodo". return_resolved_name: If True, return (model, resolved_short_key). If False, return only the model. text_encoder: Pre-built text encoder to reuse. When provided, skips text encoder selection/instantiation entirely. text_encoder_fp32: If True, uses fp32 for the text encoder rather than default bfloat16. Returns: Loaded model in eval mode, or (model, resolved short key) if return_resolved_name is True. Raises: ValueError: If modelname is not in AVAILABLE_MODELS and cannot be resolved. FileNotFoundError: If config.yaml is missing in the checkpoint folder. """ if modelname is None: modelname = DEFAULT_MODEL if modelname not in AVAILABLE_MODELS: if default_family is not None: modelname = resolve_model_name(modelname, default_family) else: raise ValueError( f"""The model is not recognized. Please choose between: {AVAILABLE_MODELS}""" ) resolved_modelname = modelname # In case, we specify a custom checkpoint directory configured_checkpoint_dir = get_env_var("CHECKPOINT_DIR") if configured_checkpoint_dir: print(f"CHECKPOINT_DIR is set to {configured_checkpoint_dir}, checking the local cache...") # Checkpoint folders are named by display name (e.g. Kimodo-SOMA-RP-v1) info = get_model_info(modelname) checkpoint_folder_name = info.display_name if info is not None else modelname model_path = Path(configured_checkpoint_dir) / checkpoint_folder_name if not model_path.exists() and modelname != checkpoint_folder_name: # Fallback: try short_key for backward compatibility model_path = Path(configured_checkpoint_dir) / modelname if not model_path.exists(): print(f"Model folder not found at '{model_path}', downloading it from Hugging Face...") model_path = _resolve_hf_model_path(modelname) else: # Otherwise, we load the model from the local cache or download it from Hugging Face. model_path = _resolve_hf_model_path(modelname) model_config_path = model_path / "config.yaml" if not model_config_path.exists(): raise FileNotFoundError(f"The model checkpoint folder exists but config.yaml is missing: {model_config_path}") model_conf = OmegaConf.load(model_config_path) if modelname in TMR_MODELS: # Same process at the moment for TMR and Kimodo pass if text_encoder is not None: runtime_conf = OmegaConf.create({"checkpoint_dir": str(model_path)}) else: text_encoder_url = get_env_var("TEXT_ENCODER_URL", DEFAULT_TEXT_ENCODER_URL) runtime_conf = OmegaConf.create( { "checkpoint_dir": str(model_path), "text_encoder": _select_text_encoder_conf(text_encoder_url, text_encoder_fp32), } ) model_cfg = OmegaConf.to_container(OmegaConf.merge(model_conf, runtime_conf), resolve=True) model_cfg.pop("checkpoint_dir", None) if text_encoder is not None: # Prevent Hydra from instantiating a new text encoder; pass None so # Kimodo.__init__ receives a placeholder we replace immediately after. model_cfg["text_encoder"] = None model = instantiate_from_dict(model_cfg, overrides={"device": device}) if text_encoder is not None: model.text_encoder = text_encoder if eval_mode: model = model.eval() if return_resolved_name: return model, resolved_modelname return model