Skip to content

vllm_omni.diffusion.models.ernie_image.pipeline_ernie_image

logger module-attribute 

logger = init_logger(__name__)

ErnieImagePipeline

Bases: Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin

current_timestep property 

current_timestep

default_sample_size instance-attribute 

default_sample_size = 128

do_classifier_free_guidance property 

do_classifier_free_guidance

guidance_scale property 

guidance_scale

image_processor instance-attribute 

image_processor = VaeImageProcessor(
    vae_scale_factor=vae_scale_factor
)

interrupt property 

interrupt

is_distilled instance-attribute 

is_distilled = is_distilled

num_timesteps property 

num_timesteps

od_config instance-attribute 

od_config = od_config

pe_model instance-attribute 

pe_model = to(_execution_device)

pe_tokenizer instance-attribute 

pe_tokenizer = from_pretrained(
    pe_base_path,
    subfolder="pe_tokenizer",
    local_files_only=True,
    trust_remote_code=True,
    use_fast=False,
)

scheduler instance-attribute 

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

support_image_input class-attribute instance-attribute 

support_image_input = False

text_encoder instance-attribute 

text_encoder = to(_execution_device)

tokenizer instance-attribute 

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

tokenizer_max_length instance-attribute 

tokenizer_max_length = 512

transformer instance-attribute 

transformer = ErnieImageTransformer2DModel(
    quant_config=quantization_config, **transformer_kwargs
)

use_pe instance-attribute 

use_pe = True

vae instance-attribute 

vae = to(_execution_device)

vae_scale_factor instance-attribute 

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

weights_sources instance-attribute 

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

check_inputs 

check_inputs(
    prompt,
    height,
    width,
    prompt_embeds=None,
    callback_on_step_end_tensor_inputs=None,
    guidance_scale=None,
)

encode_prompt 

encode_prompt(
    prompt: str | list[str],
    device: device,
    num_images_per_prompt: int = 1,
    width: int = 1024,
    height: int = 1024,
    apply_pe: bool = True,
) -> list[Tensor]

forward 

forward(
    req: OmniDiffusionRequest,
    prompt: str | list[str] | None = None,
    negative_prompt: str | list[str] | None = "",
    height: int = 1024,
    width: int = 1024,
    num_inference_steps: int = 50,
    guidance_scale: float = 4.0,
    num_images_per_prompt: int = 1,
    generator: Generator | None = None,
    latents: Tensor | None = None,
    prompt_embeds: list[FloatTensor] | None = None,
    negative_prompt_embeds: list[FloatTensor] | None = None,
    output_type: str = "pil",
    return_dict: bool = True,
    callback_on_step_end: Callable[[int, int, dict], None]
    | None = None,
    callback_on_step_end_tensor_inputs: list[str] = [
        "latents"
    ],
) -> DiffusionOutput

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

get_ernie_image_post_process_func 

get_ernie_image_post_process_func(
    od_config: OmniDiffusionConfig,
)