Skip to navigation

nemo_automodel.components.models.praxis.hub

View as MarkdownOpen in Claude

The public weights by their Hugging Face repos: the checkpoint and the frozen encoders of the reference package.

Module Contents

Functions

NameDescription
load_checkpoint_tensorsRead every shard that model.safetensors.index.json under root names.
load_text_encoderThe Qwen3.5 text encoder of the reference package.
load_vaeThe Wan2.2 VAE of the reference package.
resolve_repoLocal path of path: itself when it exists, else the Hugging Face repo path in the Hub cache.

Data

WAN_VAE_FILENAME

API

nemo_automodel.components.models.praxis.hub.load_checkpoint_tensors(
root: str | pathlib.Path,
device: str | torch.device = 'cpu'
) -> dict[str, torch.Tensor]

Read every shard that model.safetensors.index.json under root names.

Parameters:

root
str | Path

Checkpoint directory.

device
str | torch.deviceDefaults to 'cpu'

Device of the loaded tensors.

Returns: dict[str, Tensor]

The checkpoint tensors under their checkpoint keys.

nemo_automodel.components.models.praxis.hub.load_text_encoder(
path: str | pathlib.Path,
device: str | torch.device = 'cuda'
) -> praxis.text_encoder.TextEncoder

The Qwen3.5 text encoder of the reference package.

Parameters:

path
str | Path

Directory in the Hugging Face layout, or the repo id (Qwen/Qwen3.5-2B).

device
str | torch.deviceDefaults to 'cuda'

Device of the loaded model.

nemo_automodel.components.models.praxis.hub.load_vae(
path: str | pathlib.Path,
device: str | torch.device = 'cuda',
filename: str = WAN_VAE_FILENAME
) -> praxis.vae.Autoencoder

The Wan2.2 VAE of the reference package.

Parameters:

path
str | Path

The .pth weights file, or the Hugging Face repo that holds it (Wan-AI/Wan2.2-TI2V-5B).

device
str | torch.deviceDefaults to 'cuda'

Device of the loaded model.

filename
strDefaults to WAN_VAE_FILENAME

Name of the weights file inside the repo.

nemo_automodel.components.models.praxis.hub.resolve_repo(
path: str | pathlib.Path,
filename: str | None = None
) -> pathlib.Path

Local path of path: itself when it exists, else the Hugging Face repo path in the Hub cache.

Parameters:

path
str | Path

Local file or directory, or a Hugging Face repo id such as runwayml/praxis-1-droid.

filename
str | NoneDefaults to None

With a repo id, the single file to download instead of the whole repo.

Returns: Path

The local file or directory.

nemo_automodel.components.models.praxis.hub.WAN_VAE_FILENAME = 'Wan2.2_VAE.pth'