vllm_omni.diffusion.models.longcat_image ¶

default_sample_size = 128

device = get_local_device()

do_classifier_free_guidance

od_config = od_config

prompt_template_encode_prefix = "<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|im_end|>\n<|im_start|>user\n"

prompt_template_encode_suffix = (
    "<|im_end|>\n<|im_start|>assistant\n"
)

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

text_encoder = from_pretrained_with_prefetch(
    Qwen2_5_VLForConditionalGeneration.from_pretrained,
    model,
    subfolder="text_encoder",
    prefetch_list=longcat_subfolders,
    local_files_only=local_files_only,
)

text_processor = Qwen2VLProcessor.from_pretrained(
    model,
    subfolder="tokenizer",
    local_files_only=local_files_only,
)

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

tokenizer_max_length = 512

transformer = LongCatImageTransformer2DModel(
    od_config=od_config
)

vae = from_pretrained_with_prefetch(
    AutoencoderKL.from_pretrained,
    model,
    subfolder="vae",
    prefetch_list=longcat_subfolders,
    local_files_only=local_files_only,
).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,
    )
]

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

gradient_checkpointing = False

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 = in_channels

parallel_config = od_config.parallel_config

pooled_projection_dim = pooled_projection_dim

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

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

rope_preparer = RoPEPreparer(self.pos_embed)

single_transformer_blocks = nn.ModuleList(
    [
        (
            LongCatImageSingleTransformerBlock(
                parallel_config=self.parallel_config,
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
            )
        )
        for i in (range(num_single_layers))
    ]
)

time_embed = LongCatImageTimestepEmbeddings(
    embedding_dim=self.inner_dim
)

transformer_blocks = nn.ModuleList(
    [
        (
            LongCatImageTransformerBlock(
                parallel_config=self.parallel_config,
                dim=self.inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
            )
        )
        for i in (range(num_layers))
    ]
)

use_checkpoint = [True] * num_layers

use_single_checkpoint = [True] * num_single_layers

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