vllm_omni.diffusion.models.ovis_image ¶

current_timestep

default_sample_size = 128

guidance_scale

interrupt

joint_attention_kwargs

num_timesteps

od_config = od_config

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

supports_request_batch = False

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

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 = Qwen2TokenizerFast.from_pretrained(
    model,
    subfolder="tokenizer",
    local_files_only=local_files_only,
)

tokenizer_max_length = 256 + self.user_prompt_begin_id

transformer = OvisImageTransformer2DModel(
    od_config=od_config
)

user_prompt_begin_id = 28

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 = (
    2 ** (len(self.vae.config.block_out_channels) - 1)
    if getattr(self, "vae", None)
    else 8
)

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

context_embedder = nn.Linear(
    joint_attention_dim, self.inner_dim
)

context_embedder_norm = RMSNorm(
    joint_attention_dim, eps=1e-06
)

in_channels = in_channels

inner_dim = num_attention_heads * attention_head_dim

norm_out = AdaLayerNormContinuous(
    self.inner_dim,
    self.inner_dim,
    elementwise_affine=False,
    eps=1e-06,
)

out_channels = out_channels or in_channels

pos_embed = OvisImagePosEmbed(
    theta=10000, axes_dim=axes_dims_rope
)

proj_out = nn.Linear(
    self.inner_dim,
    patch_size * patch_size * self.out_channels,
    bias=True,
)

single_transformer_blocks = nn.ModuleList(
    [
        (
            OvisImageSingleTransformerBlock(
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
            )
        )
        for _ in (range(num_single_layers))
    ]
)

time_proj = Timesteps(
    num_channels=256,
    flip_sin_to_cos=True,
    downscale_freq_shift=0,
)

timestep_embedder = TimestepEmbedding(
    in_channels=256, time_embed_dim=self.inner_dim
)

transformer_blocks = nn.ModuleList(
    [
        (
            OvisImageTransformerBlock(
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
            )
        )
        for _ in (range(num_layers))
    ]
)

x_embedder = nn.Linear(in_channels, self.inner_dim)