Skip to content

vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image

logger module-attribute

logger = init_logger(__name__)

OvisImagePipeline

Bases: Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery

current_timestep property

current_timestep

default_sample_size instance-attribute

default_sample_size = 128

guidance_scale property

guidance_scale

interrupt property

interrupt

joint_attention_kwargs property

joint_attention_kwargs

num_timesteps property

num_timesteps

od_config instance-attribute

od_config = od_config

scheduler instance-attribute

scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
    model,
    subfolder="scheduler",
    local_files_only=local_files_only,
)

supports_request_batch class-attribute instance-attribute

supports_request_batch = False

system_prompt instance-attribute

system_prompt = "Describe the image by detailing the color, quantity, text, shape, size, texture, spatial\n        relationships of the objects and background: "

text_encoder instance-attribute

text_encoder = from_pretrained_with_prefetch(
    Qwen3Model.from_pretrained,
    model,
    subfolder="text_encoder",
    prefetch_list=ovis_subfolders,
    local_files_only=local_files_only,
    torch_dtype=od_config.dtype,
)

tokenizer instance-attribute

tokenizer = Qwen2TokenizerFast.from_pretrained(
    model,
    subfolder="tokenizer",
    local_files_only=local_files_only,
)

tokenizer_max_length instance-attribute

tokenizer_max_length = 256 + self.user_prompt_begin_id

transformer instance-attribute

transformer = OvisImageTransformer2DModel(
    od_config=od_config
)

user_prompt_begin_id instance-attribute

user_prompt_begin_id = 28

vae instance-attribute

vae = from_pretrained_with_prefetch(
    AutoencoderKL.from_pretrained,
    model,
    subfolder="vae",
    prefetch_list=ovis_subfolders,
    local_files_only=local_files_only,
).to(self._execution_device)

vae_scale_factor instance-attribute

vae_scale_factor = (
    2 ** (len(self.vae.config.block_out_channels) - 1)
    if getattr(self, "vae", None)
    else 8
)

weights_sources instance-attribute

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

check_inputs

check_inputs(
    prompt,
    height,
    width,
    negative_prompt=None,
    prompt_embeds=None,
    negative_prompt_embeds=None,
    callback_on_step_end_tensor_inputs=None,
    max_sequence_length=None,
)

diffuse

diffuse(
    latents: Tensor,
    timesteps: Tensor,
    prompt_embeds: Tensor,
    negative_prompt_embeds: Tensor,
    text_ids: Tensor,
    negative_text_ids: Tensor,
    latent_image_ids: Tensor,
    do_true_cfg: bool,
    guidance_scale: float,
    cfg_normalize: bool = False,
) -> Tensor

Diffusion loop with optional classifier-free guidance.

Parameters:

Name Type Description Default
latents Tensor

Noise latents to denoise

required
timesteps Tensor

Diffusion timesteps

required
prompt_embeds Tensor

Positive prompt embeddings

required
negative_prompt_embeds Tensor

Negative prompt embeddings

required
text_ids Tensor

Position IDs for positive text

required
negative_text_ids Tensor

Position IDs for negative text

required
latent_image_ids Tensor

Position IDs for image latents

required
do_true_cfg bool

Whether to apply CFG

required
guidance_scale float

CFG scale factor

required
cfg_normalize bool

Whether to normalize CFG output (default: False)

False

Returns:

Type Description
Tensor

Denoised latents

encode_prompt

encode_prompt(
    prompt: str | list[str],
    device: device | None = None,
    num_images_per_prompt: int = 1,
    prompt_embeds: FloatTensor | None = None,
)

Parameters:

Name Type Description Default
prompt `str` or `list[str]`, *optional*

prompt to be encoded

required
device device | None

(torch.device, optional): torch.device

None
num_images_per_prompt int

(int): number of images that should be generated per prompt

1
prompt_embeds FloatTensor | None

(torch.FloatTensor, optional): Pre-generated text embeddings. Can be used to easily tweak text inputs, e.g. prompt weighting. If not provided, text embeddings will be generated from prompt input argument.

None

forward

load_weights

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

prepare_latents

prepare_latents(
    batch_size,
    num_channel_latents,
    height,
    width,
    dtype,
    device,
    generator,
    latents=None,
)

prepare_timesteps

prepare_timesteps(
    num_inference_steps, sigmas, image_seq_len
)

calculate_shift

calculate_shift(
    image_seq_len,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.15,
)

get_ovis_image_post_process_func

get_ovis_image_post_process_func(
    od_config: OmniDiffusionConfig,
)

retrieve_timesteps

retrieve_timesteps(
    scheduler,
    num_inference_steps: int | None = None,
    device: str | device | None = None,
    timesteps: list[int] | None = None,
    sigmas: list[float] | None = None,
    **kwargs,
) -> tuple[Tensor, int]

Calls the scheduler's set_timesteps method and retrieves timetemps from the scheduler after the call. Handles custom timeteps. Any kwargs will be supplied to scheduler.set_timeteps.

Parameters:

Name Type Description Default
scheduler `SchedulerMixin`

The scheduler to get timesteps from.

required
num_inference_steps `int`, *optional*

The number of diffusion steps used when generating samples with a pre-trained model. If used, timesteps must be None.

None
device `str` or `torch.device`, *optional*

The device to which the timesteps should be moved to. If None, the timesteps are not moved.

None
timesteps `list[int]`, *optional*

Custom timesteps used to override the timestep spacing strategy of the scheduler. If timesteps is passed, num_inference_steps and sigmas must be None.

None
sigmas `list[float]`, *optional*

Custom sigmas used to override the timestep spacing strategy of the scheduler. If sigmas is passed, num_inference_steps and timesteps must be None.

None

Returns:

Type Description
Tensor

Tuple[torch.Tensor, int]: A tuple where the first element is the timestep schedule from the scheduler and the

int

second element is the number of inference steps.