vllm_omni.diffusion.models.dreamzero.image_encoder ¶

model = _DreamZeroCLIPContainer()

transforms = Compose(
    [
        Normalize(
            mean=[0.48145466, 0.4578275, 0.40821073],
            std=[0.26862954, 0.26130258, 0.27577711],
        )
    ]
)

attn = DreamZeroVisionSelfAttention(
    dim, num_heads, proj_dropout=proj_dropout
)

mlp = Sequential(
    Linear(dim, hidden_dim),
    GELU(),
    Linear(hidden_dim, dim),
    Dropout(proj_dropout),
)

norm1 = DreamZeroLayerNorm(dim, eps=norm_eps)

norm2 = DreamZeroLayerNorm(dim, eps=norm_eps)

post_norm = post_norm

dim = dim

head_dim = dim // num_heads

num_heads = num_heads

proj = Linear(dim, dim)

proj_dropout = proj_dropout

to_qkv = Linear(dim, dim * 3)

cls_embedding = Parameter(gain * randn(1, 1, dim))

dim = dim

dropout = Dropout(embedding_dropout)

head = Parameter(gain * randn(dim, out_dim))

image_size = image_size

num_heads = num_heads

num_layers = num_layers

num_patches = (image_size // patch_size) ** 2

patch_embedding = Conv2d(
    3,
    dim,
    kernel_size=patch_size,
    stride=patch_size,
    bias=not pre_norm,
)

patch_size = patch_size

pool_type = pool_type

pos_embedding = Parameter(
    gain * randn(1, num_patches + 1, dim)
)

post_norm = DreamZeroLayerNorm(dim, eps=norm_eps)

pre_norm = (
    DreamZeroLayerNorm(dim, eps=norm_eps)
    if pre_norm
    else None
)

transformer = Sequential(
    *[
        (
            DreamZeroVisionAttentionBlock(
                dim=dim,
                mlp_ratio=mlp_ratio,
                num_heads=num_heads,
                post_norm=post_norm,
                activation=activation,
                proj_dropout=proj_dropout,
                norm_eps=norm_eps,
            )
        )
        for _ in (range(num_layers))
    ]
)