Skip to content

vllm_omni.diffusion.model_loader.diffusers_loader

DIFFUSION_MODEL_WEIGHTS_INDEX module-attribute 

DIFFUSION_MODEL_WEIGHTS_INDEX = (
    "diffusion_pytorch_model.safetensors.index.json"
)

INDEX_FILES module-attribute 

INDEX_FILES = [
    DIFFUSION_MODEL_WEIGHTS_INDEX,
    TRANSFORMER_WEIGHTS_INDEX,
]

MODEL_INDEX module-attribute 

MODEL_INDEX = 'model_index.json'

TRANSFORMER_WEIGHTS_INDEX module-attribute 

TRANSFORMER_WEIGHTS_INDEX = 'model.safetensors.index.json'

logger module-attribute 

logger = init_logger(__name__)

DiffusersPipelineLoader

Model loader that can load diffusers pipeline components from disk.

counter_after_loading_weights class-attribute instance-attribute 

counter_after_loading_weights: float = 0.0

counter_before_loading_weights class-attribute instance-attribute 

counter_before_loading_weights: float = 0.0

load_config instance-attribute 

load_config = load_config

od_config instance-attribute 

od_config = od_config

ComponentSource dataclass

A source for weights.

allow_patterns_overrides class-attribute instance-attribute 

allow_patterns_overrides: list[str] | None = None

If defined, weights will load exclusively using these patterns.

fall_back_to_pt class-attribute instance-attribute 

fall_back_to_pt: bool = True

Whether .pt weights can be used.

model_or_path instance-attribute 

model_or_path: str

The model ID or path.

prefix class-attribute instance-attribute 

prefix: str = ''

A prefix to prepend to all weights.

revision instance-attribute 

revision: str | None

The optional model revision.

subfolder instance-attribute 

subfolder: str | None

The subfolder inside the model repo.

download_model 

download_model(model_config: ModelConfig) -> None

get_all_weights 

get_all_weights(
    model: Module,
) -> Generator[tuple[str, Tensor], None, None]

load_model 

load_model(
    od_config: OmniDiffusionConfig,
    load_device: str,
    load_format: str | None = "default",
    custom_pipeline_name: str | type[Module] | None = None,
    device: device | None = None,
) -> Module

Load a model with the given configurations.

load_weights 

load_weights(model: Module) -> None