omnihuman.utils.io module

IO Module

This module provides utility functions for reading and writing files.

Functions:

  • read_frames: Read frames from image or video file as 4D tensor (n_frames, n_channels, height, width).

  • fetch_pretrained_weights: Downloads & reads specific tensors from a Hugging Face Hub repository.

omnihuman.utils.io.fetch_pretrained_weights(repo_id: str, weight_name_to_file_name: Dict[str, str]) → Dict[str, Tensor][source][source]

Downloads specific files from a Hugging Face Hub repository and loads specific tensors from them.

Parameters:

  • repo_id (str) – Model repository name on Hugging Face Hub (e.g. “organization/their-awesome-model”).

  • weight_name_to_file_name (Dict[str, str]) – mapping of layer name to the shard file that contains its weights (Explore the model.safetensors.index.json in the HF model repo for names)

Returns:

mapping from layer name to its weights tensor

Return type:

Dict[str, torch.Tensor]

omnihuman.utils.io.read_frames(path: str) → Tensor[source][source]

Read frames from image or video file as 4D tensor (n_frames, n_channels, height, width).

Parameters:

path (str) – Where the image or video file is located.

Raises:

ValueError – If the file type is neither image nor video.

Returns:

Frames as 4d torch tensor of shape (n_frames, n_channels, height, width).

Return type:

torch.Tensor