Skip to content

vllm_omni.diffusion.models.hidream_o1_image

Modules:

Name Description
hidream_o1_image_transformer
pipeline_hidream_o1_image

HiDreamO1ImagePipeline

Bases: Module, DiffusionPipelineProfilerMixin, ProgressBarMixin

device instance-attribute 

device = get_local_device()

dtype instance-attribute 

dtype = (
    od_config.dtype
    if od_config.dtype is not None
    else torch.bfloat16
)

dummy_run_num_frames class-attribute 

dummy_run_num_frames: int = 0

model_dir instance-attribute 

model_dir = od_config.model

od_config instance-attribute 

od_config = od_config

prefix instance-attribute 

prefix = prefix

supports_request_batch class-attribute instance-attribute 

supports_request_batch = False

weights_sources instance-attribute 

weights_sources = [
    DiffusersPipelineLoader.ComponentSource(
        model_or_path=od_config.model,
        subfolder=None,
        revision=od_config.revision,
        prefix="model.",
        fall_back_to_pt=True,
    )
]

forward 

forward(req: DiffusionRequestBatch) -> DiffusionOutput

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

HiDreamO1ImageTransformer

Bases: Module

config instance-attribute 

config = config

model instance-attribute 

model = HiDreamO1ImageModel(
    config,
    quant_config=quant_config,
    prefix=_prefix(prefix, "model"),
)

quant_config instance-attribute 

quant_config = quant_config

supports_piecewise_attention property 

supports_piecewise_attention: bool

forward 

forward(
    input_ids: Tensor,
    position_ids: Tensor,
    vinputs: Tensor,
    timestep: Tensor,
    attention_mask: Tensor | None,
    full_attn_spans: list[list[tuple[int, int]]]
    | None = None,
    **_: object,
) -> HiDreamO1ImageOutput

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

get_hidream_o1_image_post_process_func 

get_hidream_o1_image_post_process_func(
    od_config: OmniDiffusionConfig,
)