vllm_omni.diffusion.models.sd3 ¶

attention_head_dim = attention_head_dim

caption_projection_dim = caption_projection_dim

context_embedder = ReplicatedLinear(
    joint_attention_dim, caption_projection_dim
)

dual_attention_layers = (
    dual_attention_layers
    if hasattr(model_config, "dual_attention_layers")
    else ()
)

in_channels = in_channels

inner_dim = num_attention_heads * attention_head_dim

joint_attention_dim = joint_attention_dim

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

num_attention_heads = num_attention_heads

num_layers = num_layers

out_channels = out_channels

parallel_config = parallel_config

patch_size = patch_size

pooled_projection_dim = pooled_projection_dim

pos_embed = PatchEmbed(
    height=sample_size,
    width=sample_size,
    patch_size=patch_size,
    in_channels=in_channels,
    embed_dim=inner_dim,
    pos_embed_max_size=pos_embed_max_size,
)

pos_embed_max_size = pos_embed_max_size

proj_out = ReplicatedLinear(
    inner_dim,
    patch_size * patch_size * out_channels,
    bias=True,
)

qk_norm = (
    qk_norm if hasattr(model_config, "qk_norm") else ""
)

sample_size = sample_size

time_text_embed = CombinedTimestepTextProjEmbeddings(
    embedding_dim=inner_dim,
    pooled_projection_dim=pooled_projection_dim,
)

transformer_blocks = ModuleList(
    [
        (
            SD3TransformerBlock(
                dim=inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                context_pre_only=i == num_layers - 1,
                qk_norm=qk_norm,
                use_dual_attention=True
                if i in dual_attention_layers
                else False,
            )
        )
        for i in (range(num_layers))
    ]
)

current_timestep

default_sample_size = 128

device = get_local_device()

guidance_scale

image_processor = VaeImageProcessor(
    vae_scale_factor=vae_scale_factor
)

interrupt

num_timesteps

od_config = od_config

output_type = output_type

patch_size = 2

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

text_encoder = from_pretrained_with_prefetch(
    from_pretrained,
    model,
    subfolder="text_encoder",
    prefetch_list=sd3_subfolders,
    local_files_only=local_files_only,
    torch_dtype=dtype,
)

text_encoder_2 = from_pretrained_with_prefetch(
    from_pretrained,
    model,
    subfolder="text_encoder_2",
    prefetch_list=sd3_subfolders,
    local_files_only=local_files_only,
    torch_dtype=dtype,
)

text_encoder_3 = from_pretrained_with_prefetch(
    from_pretrained,
    model,
    subfolder="text_encoder_3",
    prefetch_list=sd3_subfolders,
    local_files_only=local_files_only,
    torch_dtype=dtype,
)

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

tokenizer_2 = from_pretrained(
    model,
    subfolder="tokenizer_2",
    local_files_only=local_files_only,
)

tokenizer_3 = from_pretrained(
    model,
    subfolder="tokenizer_3",
    local_files_only=local_files_only,
)

tokenizer_max_length = (
    model_max_length
    if hasattr(self, "tokenizer") and tokenizer is not None
    else 77
)

transformer = SD3Transformer2DModel(od_config=od_config)

vae = to(device)

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

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