vllm_omni.diffusion.models.sd3 ¶

attention_head_dim = model_config.attention_head_dim

caption_projection_dim = model_config.caption_projection_dim

context_embedder = ReplicatedLinear(
    self.joint_attention_dim, self.caption_projection_dim
)

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

in_channels = model_config.in_channels

inner_dim = (
    model_config.num_attention_heads
    * model_config.attention_head_dim
)

joint_attention_dim = model_config.joint_attention_dim

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

num_attention_heads = model_config.num_attention_heads

num_layers = model_config.num_layers

out_channels = model_config.out_channels

parallel_config = od_config.parallel_config

patch_size = model_config.patch_size

pooled_projection_dim = model_config.pooled_projection_dim

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

pos_embed_max_size = model_config.pos_embed_max_size

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

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

sample_size = model_config.sample_size

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

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

current_timestep

default_sample_size = 128

device = get_local_device()

guidance_scale

image_processor = VaeImageProcessor(
    vae_scale_factor=self.vae_scale_factor
)

interrupt

num_timesteps

od_config = od_config

output_type = self.od_config.output_type

patch_size = 2

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

text_encoder = from_pretrained_with_prefetch(
    CLIPTextModelWithProjection.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(
    CLIPTextModelWithProjection.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(
    T5EncoderModel.from_pretrained,
    model,
    subfolder="text_encoder_3",
    prefetch_list=sd3_subfolders,
    local_files_only=local_files_only,
    torch_dtype=dtype,
)

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

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

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

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

transformer = SD3Transformer2DModel(od_config=od_config)

vae = from_pretrained_with_prefetch(
    DistributedAutoencoderKL.from_pretrained,
    model,
    subfolder="vae",
    prefetch_list=sd3_subfolders,
    local_files_only=local_files_only,
    torch_dtype=dtype,
).to(self.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,
    )
]