Skip to content

vllm_omni.diffusion.models.sdxl.pipeline_sdxl

logger module-attribute

logger = logging.getLogger(__name__)

StableDiffusionXLPipeline

Bases: Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery

current_timestep property

current_timestep

default_sample_size instance-attribute

default_sample_size = 128

device instance-attribute

device = get_local_device()

guidance_scale property

guidance_scale

image_processor instance-attribute

image_processor = VaeImageProcessor(
    vae_scale_factor=self.vae_scale_factor
)

interrupt property

interrupt

num_timesteps property

num_timesteps

od_config instance-attribute

od_config = od_config

output_type instance-attribute

output_type = self.od_config.output_type

scheduler instance-attribute

scheduler = EulerDiscreteScheduler.from_pretrained(
    model,
    subfolder="scheduler",
    local_files_only=local_files_only,
    torch_dtype=torch.float32,
)

text_encoder instance-attribute

text_encoder = from_pretrained_with_prefetch(
    CLIPTextModel.from_pretrained,
    model,
    subfolder="text_encoder",
    prefetch_list=sdxl_subfolders,
    local_files_only=local_files_only,
    torch_dtype=dtype,
)

text_encoder_2 instance-attribute

text_encoder_2 = from_pretrained_with_prefetch(
    CLIPTextModelWithProjection.from_pretrained,
    model,
    subfolder="text_encoder_2",
    prefetch_list=sdxl_subfolders,
    local_files_only=local_files_only,
    torch_dtype=dtype,
)

tokenizer instance-attribute

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

tokenizer_2 instance-attribute

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

tokenizer_max_length instance-attribute

tokenizer_max_length = 77

unet instance-attribute

unet = SDXLUNet2DConditionModel(od_config=od_config)

vae instance-attribute

vae = from_pretrained_with_prefetch(
    DistributedAutoencoderKL.from_pretrained,
    model,
    subfolder="vae",
    prefetch_list=sdxl_subfolders,
    local_files_only=local_files_only,
    torch_dtype=dtype,
).to(self.device)

vae_scale_factor instance-attribute

vae_scale_factor = 8

weights_sources instance-attribute

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

diffuse

diffuse(
    latents: Tensor,
    timesteps: Tensor,
    prompt_embeds: Tensor,
    added_cond_kwargs: dict,
    negative_prompt_embeds: Tensor | None,
    negative_added_cond_kwargs: dict | None,
    do_cfg: bool,
    guidance_scale: float,
) -> Tensor

encode_prompt

encode_prompt(
    prompt: str | list[str], num_images_per_prompt: int = 1
) -> tuple[Tensor, Tensor]

forward

forward(
    req: OmniDiffusionRequest,
    prompt: str | list[str] = "",
    negative_prompt: str | list[str] = "",
    height: int | None = None,
    width: int | None = None,
    num_inference_steps: int = 50,
    num_images_per_prompt: int = 1,
    generator: Generator | list[Generator] | None = None,
    latents: Tensor | None = None,
) -> DiffusionOutput

load_weights

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

predict_noise

predict_noise(**kwargs)

prepare_latents

prepare_latents(
    batch_size: int,
    num_channels_latents: int,
    height: int,
    width: int,
    generator: Generator | list[Generator] | None,
    latents: Tensor | None = None,
) -> Tensor

get_sdxl_image_post_process_func

get_sdxl_image_post_process_func(
    od_config: OmniDiffusionConfig,
)