Skip to content

vllm_omni.diffusion.models.z_image.pipeline_z_image

logger module-attribute 

logger = logging.get_logger(__name__)

ZImagePipeline

Bases: Module, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery

do_classifier_free_guidance property 

do_classifier_free_guidance

guidance_scale property 

guidance_scale

image_processor instance-attribute 

image_processor = VaeImageProcessor(
    vae_scale_factor=self.vae_scale_factor * 2,
    do_convert_rgb=True,
)

interrupt property 

interrupt

joint_attention_kwargs property 

joint_attention_kwargs

num_timesteps property 

num_timesteps

od_config instance-attribute 

od_config = od_config

scheduler instance-attribute 

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

supports_request_batch class-attribute instance-attribute 

supports_request_batch = False

text_encoder instance-attribute 

text_encoder = create_transformers_model(
    AutoModelForCausalLM,
    od_config,
    hf_config=text_encoder_config,
).to(self._execution_device)

tokenizer instance-attribute 

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

transformer instance-attribute 

transformer = ZImageTransformer2DModel(
    quant_config=od_config.quantization_config
)

vae instance-attribute 

vae = DistributedAutoencoderKL.from_config(vae_config).to(
    self._execution_device
)

vae_scale_factor instance-attribute 

vae_scale_factor = (
    2 ** (len(self.vae.config.block_out_channels) - 1)
    if hasattr(self, "vae") and self.vae is not None
    else 8
)

weights_sources instance-attribute 

weights_sources = [
    DiffusersPipelineLoader.ComponentSource(
        model_or_path=od_config.model,
        subfolder="text_encoder",
        revision=od_config.revision,
        prefix="text_encoder.",
    ),
    DiffusersPipelineLoader.ComponentSource(
        model_or_path=od_config.model,
        subfolder="transformer",
        revision=od_config.revision,
        prefix="transformer.",
        fall_back_to_pt=True,
    ),
    DiffusersPipelineLoader.ComponentSource(
        model_or_path=od_config.model,
        subfolder="vae",
        revision=od_config.revision,
        prefix="vae.",
    ),
]

encode_prompt 

encode_prompt(
    prompt: str | list[str],
    device: device | None = None,
    do_classifier_free_guidance: bool = True,
    negative_prompt: str | list[str] | None = None,
    prompt_embeds: list[FloatTensor] | None = None,
    negative_prompt_embeds: FloatTensor | None = None,
    max_sequence_length: int = 512,
)

forward 

forward(req: DiffusionRequestBatch) -> DiffusionOutput

get_timesteps 

get_timesteps(num_inference_steps, strength, device)

load_weights 

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

prepare_latents 

prepare_latents(
    batch_size,
    num_channels_latents,
    height,
    width,
    dtype,
    device,
    generator,
    latents=None,
    image=None,
    timestep=None,
)

calculate_shift 

calculate_shift(
    image_seq_len,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.15,
)

get_post_process_func 

get_post_process_func(od_config: OmniDiffusionConfig)

retrieve_latents 

retrieve_latents(
    encoder_output: Tensor,
    generator: Generator | None = None,
    sample_mode: str = "sample",
)

retrieve_timesteps 

retrieve_timesteps(
    scheduler,
    num_inference_steps: int | None = None,
    device: str | device | None = None,
    timesteps: list[int] | None = None,
    sigmas: list[float] | None = None,
    **kwargs,
) -> tuple[Tensor, int]

Calls the scheduler's set_timesteps method and retrieves timesteps from the scheduler after the call. Handles custom timesteps. Any kwargs will be supplied to scheduler.set_timesteps.

Parameters:

Name Type Description Default
scheduler `SchedulerMixin`

The scheduler to get timesteps from.

 required
num_inference_steps `int`

The number of diffusion steps used when generating samples with a pre-trained model. If used, timesteps must be None.

 None
device `str` or `torch.device`, *optional*

The device to which the timesteps should be moved to. If None, the timesteps are not moved.

 None
timesteps `list[int]`, *optional*

Custom timesteps used to override the timestep spacing strategy of the scheduler. If timesteps is passed, num_inference_steps and sigmas must be None.

 None
sigmas `list[float]`, *optional*

Custom sigmas used to override the timestep spacing strategy of the scheduler. If sigmas is passed, num_inference_steps and timesteps must be None.

 None

Returns:

Type Description
Tensor

Tuple[torch.Tensor, int]: A tuple where the first element is the timestep schedule from the scheduler and the
int

second element is the number of inference steps.