vllm_omni.diffusion.models.boogu_image ¶

SYSTEM_PROMPT_4_I2I = SYSTEM_PROMPT_4_TI2I_UNIFIED

SYSTEM_PROMPT_4_T2I = SYSTEM_PROMPT_4_T2I_UNIFIED

SYSTEM_PROMPT_4_TI2I = SYSTEM_PROMPT_4_TI2I_UNIFIED

SYSTEM_PROMPT_DROP = SYSTEM_PROMPT_4_TI2I_UNIFIED

color_format: str = 'RGB'

default_sample_size = 128

mllm = self._load_mllm(
    model,
    local_files_only,
    mllm_quant_config,
    boogu_subfolders,
)

od_config = od_config

processor = Qwen3VLProcessor.from_pretrained(
    model,
    subfolder="processor",
    local_files_only=local_files_only,
    revision=od_config.revision,
)

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

support_image_input: bool = True

supports_request_batch = True

transformer = BooguImageTransformer2DModel(
    od_config=od_config,
    quant_config=transformer_quant_config,
    prefix="transformer",
)

vae = from_pretrained_with_prefetch(
    AutoencoderKL.from_pretrained,
    model,
    subfolder="vae",
    prefetch_list=boogu_subfolders,
    local_files_only=local_files_only,
    revision=od_config.revision,
).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=od_config.revision,
        prefix="transformer.",
        fall_back_to_pt=True,
    )
]

max_pixels = max_pixels

max_side_length = max_side_length

axes_dim_rope = axes_dim_rope

axes_lens = axes_lens

context_refiner = nn.ModuleList(
    [
        BooguImageContextRefinerTransformerBlock(
            hidden_size,
            num_attention_heads,
            num_kv_heads,
            multiple_of,
            ffn_dim_multiplier,
            norm_eps,
            modulation=False,
            quant_config=quant_config,
            prefix=_join_prefix(
                prefix, f"context_refiner.{i}"
            ),
        )
        for i in range(num_refiner_layers)
    ]
)

double_stream_layers = nn.ModuleList(
    [
        BooguImageDoubleStreamTransformerBlock(
            hidden_size,
            num_attention_heads,
            num_kv_heads,
            multiple_of,
            ffn_dim_multiplier,
            norm_eps,
            modulation=True,
            quant_config=quant_config,
            prefix=_join_prefix(
                prefix, f"double_stream_layers.{i}"
            ),
        )
        for i in range(num_double_stream_layers)
    ]
)

hidden_size = hidden_size

image_index_embedding = nn.Parameter(
    torch.randn(5, hidden_size)
)

in_channels = in_channels

instruction_feature_configs = instruction_feature_configs

noise_refiner = nn.ModuleList(
    [
        BooguImageNoiseRefinerTransformerBlock(
            hidden_size,
            num_attention_heads,
            num_kv_heads,
            multiple_of,
            ffn_dim_multiplier,
            norm_eps,
            modulation=True,
            quant_config=quant_config,
            prefix=_join_prefix(
                prefix, f"noise_refiner.{i}"
            ),
        )
        for i in range(num_refiner_layers)
    ]
)

norm_out = LuminaLayerNormContinuous(
    embedding_dim=hidden_size,
    conditioning_embedding_dim=min(hidden_size, 1024),
    elementwise_affine=False,
    eps=1e-06,
    bias=True,
    out_dim=patch_size * patch_size * self.out_channels,
    quant_config=quant_config,
    prefix=_join_prefix(prefix, "norm_out"),
)

num_attention_heads = num_attention_heads

num_double_stream_layers = num_double_stream_layers

num_single_stream_layers = (
    num_layers - num_double_stream_layers
)

od_config = od_config

out_channels = out_channels or in_channels

patch_size = patch_size

preprocessed_instruction_feat_dim = (
    _cal_preprocessed_instruction_feat_dim(
        instruction_feature_configs
    )
)

ref_image_patch_embedder = nn.Linear(
    in_features=patch_size * patch_size * in_channels,
    out_features=hidden_size,
)

ref_image_refiner = nn.ModuleList(
    [
        BooguImageRefImgRefinerTransformerBlock(
            hidden_size,
            num_attention_heads,
            num_kv_heads,
            multiple_of,
            ffn_dim_multiplier,
            norm_eps,
            modulation=True,
            quant_config=quant_config,
            prefix=_join_prefix(
                prefix, f"ref_image_refiner.{i}"
            ),
        )
        for i in range(num_refiner_layers)
    ]
)

rope_embedder = BooguImageDoubleStreamRotaryPosEmbed(
    theta=10000,
    axes_dim=axes_dim_rope,
    axes_lens=axes_lens,
    patch_size=patch_size,
)

single_stream_layers = nn.ModuleList(
    [
        BooguImageSingleStreamTransformerBlock(
            hidden_size,
            num_attention_heads,
            num_kv_heads,
            multiple_of,
            ffn_dim_multiplier,
            norm_eps,
            modulation=True,
            quant_config=quant_config,
            prefix=_join_prefix(
                prefix, f"single_stream_layers.{i}"
            ),
        )
        for i in range(self.num_single_stream_layers)
    ]
)

stacked_params_mapping = list(_BOOGU_STACKED_PARAMS_MAPPING)

time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding(
    hidden_size=hidden_size,
    instruction_feat_dim=self.preprocessed_instruction_feat_dim,
    norm_eps=norm_eps,
    timestep_scale=timestep_scale,
)

x_embedder = nn.Linear(
    in_features=patch_size * patch_size * in_channels,
    out_features=hidden_size,
)

supports_request_batch = False

begin_index

order = 1

step_index

time_shift_v2_scaling_factor = (
    time_shift_v2_half_scaling_factor * 2
)

timesteps = timesteps