load_model_component

Load a specific component of a diffusion model

Description

Loads a TorchScript model component (UNet, decoder, or text encoder) from the hfhub cache or legacy model directory, downloading it first if necessary.

Usage

load_model_component(component, model_name = "sd21", device = "cpu",
                     unet_dtype_str = NULL, download = TRUE, use_native = FALSE)

Arguments

  • component: Character string, the component to load: “unet”, “decoder”, or “text_encoder”.
  • model_name: Character string, the name of the model to use.
  • device: Character string, the torch device to load the model onto (“cpu” or “cuda”).
  • unet_dtype_str: Optional; the data type for the UNet model. If NULL, defaults to float32 for CPU and float16 for CUDA.
  • download: Logical; if TRUE (default), downloads the model if it doesn’t exist locally.
  • use_native: Logical; if TRUE, uses native R torch modules instead of TorchScript. Supported for unet, decoder, text_encoder, and text_encoder2. Native modules have better GPU compatibility (especially on Blackwell/RTX 50xx).

Value

A torch model object.

Examples

unet <- load_model_component("unet", "sd21", "cpu")