Skip to content

vllm_omni.diffusion.model_loader.diffusers_loader

DIFFUSION_MODEL_BIN_WEIGHTS_INDEX module-attribute 

DIFFUSION_MODEL_BIN_WEIGHTS_INDEX = (
    "diffusion_pytorch_model.bin.index.json"
)

DIFFUSION_MODEL_WEIGHTS_INDEX module-attribute 

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

PT_INDEX_FILES module-attribute 

PT_INDEX_FILES = [DIFFUSION_MODEL_BIN_WEIGHTS_INDEX]

SAFETENSORS_INDEX_FILES module-attribute 

SAFETENSORS_INDEX_FILES = [
    DIFFUSION_MODEL_WEIGHTS_INDEX,
    TRANSFORMER_WEIGHTS_INDEX,
]

SHARDED_SAFETENSORS_PATTERN module-attribute 

SHARDED_SAFETENSORS_PATTERN = re.compile(
    "^(?P<family>.+)-\\d+-of-(?P<count>\\d+)\\.safetensors$"
)

TRANSFORMER_WEIGHTS_INDEX module-attribute 

TRANSFORMER_WEIGHTS_INDEX = 'model.safetensors.index.json'

logger module-attribute 

logger = init_logger(__name__)

DiffusersPipelineLoader

Bases: HWRLoaderMixin

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

host_weight_plan instance-attribute 

host_weight_plan: HostWeightPlan | None = None

load_config instance-attribute 

load_config = load_config

od_config instance-attribute 

od_config = od_config

parallel_config instance-attribute 

parallel_config = od_config.parallel_config

quant_config instance-attribute 

quant_config = od_config.quantization_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.

get_all_weights 

get_all_weights(
    model: Module,
    sources: Sequence[ComponentSource] | None = None,
) -> Generator[tuple[str, Tensor], None, None]

load_model 

load_model(
    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,
    *,
    stream_online_quant_to_cpu: bool = False,
    sources: Sequence[ComponentSource] | None = None,
    planned_weights: Iterable[str] = (),
) -> None

take_host_weight_plan 

take_host_weight_plan() -> HostWeightPlan | None

Transfer the loader-produced plan to the offload backend.

download_gguf 

download_gguf(
    repo_id: str,
    quant_type: str,
    cache_dir: str | None = None,
    revision: str | None = None,
    ignore_patterns: str | list[str] | None = None,
) -> str