Diffusers documentation
PRX
PRX
PRX 使用简化的 MMDIT 架构从文本生成高质量图像，在该架构中，文本 token 不会通过 transformer 块进行更新。它采用带有离散调度的 flow matching 以实现高效采样，并使用 Google 的 T5Gemma-2B-2B-UL2 模型进行多语言文本编码。该约 13 亿参数的 transformer 在不牺牲质量的前提下实现了快速推理。您可以选择 Flux VAE（8 倍压缩，16 个潜在通道）以平衡质量和速度，或选择 DC-AE（32 倍压缩，32 个潜在通道）以实现潜在压缩和更快的处理。
可用模型
PRX 提供多种具有不同 VAE 配置的变体，每种变体均针对特定分辨率进行了优化。基础模型在处理详细提示词方面表现出色，能够捕捉复杂的构图和细微的细节。在 Alchemist 数据集 上经过微调的模型提升了美学质量，尤其是在使用较简单的提示词时。
| 模型 | 分辨率 | 微调 | 蒸馏 | 描述 | 建议的提示 | 建议参数 | 推荐的 dtype |
|---|---|---|---|---|---|---|---|
Photoroom/prx-256-t2i | 256 | 否 | 否 | 使用 Flux VAE 在 256 分辨率下预训练的基础模型 | 在自然语言中提供详细的提示效果最佳 | 28 步，cfg=5.0 | torch.bfloat16 |
Photoroom/prx-256-t2i-sft | 512 | 是的 | 否 | 使用 Flux VAE 在 Alchemist dataset 数据集上进行微调 | 能够处理细节较少的提示 | 28 步，cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i | 512 | 否 | 否 | 使用 Flux VAE 在 512 分辨率下预训练的基础模型 | 在自然语言中提供详细的提示效果最佳 | 28 步，cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-sft | 512 | 是的 | 否 | 使用 Flux VAE 在 Alchemist dataset 数据集上进行微调 | 能够处理自然语言中细节较少的提示 | 28 步，cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-sft-distilled | 512 | 是的 | 是的 | 来自 Photoroom/prx-512-t2i-sft 的 8 步蒸馏模型 | 能够处理自然语言中细节较少的提示 | 8 步，cfg=1.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae | 512 | 否 | 否 | 使用 Deep Compression Autoencoder (DC-AE) 在 512 分辨率下预训练的基础模型 | 在自然语言中提供详细的提示效果最佳 | 28 步，cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae-sft | 512 | 是的 | 否 | 使用 Deep Compression Autoencoder (DC-AE) 在 Alchemist dataset 数据集上进行微调 | 能够处理自然语言中细节较少的提示 | 28 步，cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae-sft-distilled | 512 | 是的 | 是的 | 来自 Photoroom/prx-512-t2i-dc-ae-sft-distilled 的 8 步蒸馏模型 | 能够处理自然语言中细节较少的提示 | 8 步，cfg=1.0 | torch.bfloat16 |
有关更多信息，请参阅此集合。
加载 pipeline
使用 from_pretrained() 加载 pipeline。
from diffusers.pipelines.prx import PRXPipeline
# Load pipeline - VAE and text encoder will be loaded from HuggingFace
pipe = PRXPipeline.from_pretrained("Photoroom/prx-512-t2i-sft", dtype=torch.bfloat16)
pipe.to("cuda") # or "mps", "xpu", "cpu"
prompt = "A front-facing portrait of a lion the golden savanna at sunset."
image = pipe(prompt, num_inference_steps=28, guidance_scale=5.0).images[0]
image.save("prx_output.png")手动加载组件
单独加载各个组件，以自定义流水线，例如使用量化模型。
import torch
from diffusers.pipelines.prx import PRXPipeline
from diffusers.models import AutoencoderKL, AutoencoderDC
from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from transformers import T5GemmaModel, GemmaTokenizerFast
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as BitsAndBytesConfig
quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
# Load transformer
transformer = PRXTransformer2DModel.from_pretrained(
"checkpoints/prx-512-t2i-sft",
subfolder="transformer",
quantization_config=quant_config,
dtype=torch.bfloat16,
)
# Load scheduler
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
"checkpoints/prx-512-t2i-sft", subfolder="scheduler"
)
# Load T5Gemma text encoder
t5gemma_model = T5GemmaModel.from_pretrained("google/t5gemma-2b-2b-ul2",
quantization_config=quant_config,
dtype=torch.bfloat16)
text_encoder = t5gemma_model.encoder.to(dtype=torch.bfloat16)
tokenizer = GemmaTokenizerFast.from_pretrained("google/t5gemma-2b-2b-ul2")
tokenizer.model_max_length = 256
# Load VAE - choose either Flux VAE or DC-AE
# Flux VAE
vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev",
subfolder="vae",
quantization_config=quant_config,
dtype=torch.bfloat16)
pipe = PRXPipeline(
transformer=transformer,
scheduler=scheduler,
text_encoder=text_encoder,
tokenizer=tokenizer,
vae=vae
)
pipe.to("cuda") # or "mps", "xpu", "cpu"内存优化
对于内存受限的环境：
import torch
from diffusers.pipelines.prx import PRXPipeline
pipe = PRXPipeline.from_pretrained("Photoroom/prx-512-t2i-sft", dtype=torch.bfloat16)
pipe.enable_model_cpu_offload() # Offload components to CPU when not in use
# Or use sequential CPU offload for even lower memory
pipe.enable_sequential_cpu_offload()PRXPipeline
class diffusers.PRXPipeline
< source >( transformer: PRXTransformer2DModelscheduler: FlowMatchEulerDiscreteSchedulertext_encoder: T5GemmaEncodertokenizer: transformers.models.t5.tokenization_t5.T5Tokenizer | transformers.models.gemma.tokenization_gemma.GemmaTokenizer | transformers.models.auto.tokenization_auto.AutoTokenizervae: diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL | diffusers.models.autoencoders.autoencoder_dc.AutoencoderDC | None = Nonedefault_sample_size: int | None = 512 )
参数
- transformer (
PRXTransformer2DModel) — 用于对编码后的图像潜变量进行去噪的 PRX transformer 模型。 - scheduler (FlowMatchEulerDiscreteScheduler) —
用于与
transformer配合，对编码后的图像潜变量进行去噪的调度器。 - text_encoder (
T5GemmaEncoder) — 用于编码提示词的文本编码器模型。 - tokenizer（
T5TokenizerFast或GemmaTokenizerFast）— 文本编码器的分词器。 - vae（AutoencoderKL 或 AutoencoderDC）—— 变分自编码器（VAE）模型，用于将图像编码为潜在表示，以及从潜在表示解码回图像。 同时支持 AutoencoderKL（8 倍压缩）和 AutoencoderDC（32 倍压缩）。
使用 PRX Transformer 进行文本到图像生成的 Pipeline。
该模型继承自 DiffusionPipeline。有关该库为所有流水线实现的通用方法（例如下载或保存、在特定设备上运行等），请参阅父类文档。
__call__
< source >( prompt: str | list[str] = Nonenegative_prompt: str = ''height: int | None = Nonewidth: int | None = Nonenum_inference_steps: int = 28timesteps: list = Noneguidance_scale: float = 4.0num_images_per_prompt: int | None = 1generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneprompt_embeds: typing.Optional[torch.FloatTensor] = Nonenegative_prompt_embeds: typing.Optional[torch.FloatTensor] = Noneprompt_attention_mask: typing.Optional[torch.BoolTensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.BoolTensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Trueuse_resolution_binning: bool = Truecallback_on_step_end: typing.Optional[typing.Callable[[int, int], NoneType]] = Nonecallback_on_step_end_tensor_inputs: list = ['latents']tokenizer_max_length: int | None = Noneskip_text_cleaning: bool = False ) → [~pipelines.prx.PRXPipelineOutput] or *tuple`
参数
- prompt (str 或 list[str]，可选) — 用于引导图像生成的提示词或提示词列表。如果未定义，则必须改为传入 prompt_embeds。
- negative_prompt (str，可选，默认为 "") — 用于不引导图像生成的提示词。当未使用引导时会被忽略（即当 guidance_scale 小于 1 时会被忽略）。
- height (int，可选，默认为 self.transformer.config.sample_size * self.vae_scale_factor) — 生成图像的高度（以像素为单位）。
- width (int，可选，默认为 self.transformer.config.sample_size * self.vae_scale_factor) — 生成图像的宽度（以像素为单位）。
- num_inference_steps (int，可选，默认为 28) — 去噪步数。更多的去噪步数通常能生成更高质量的图像，但代价是推理速度变慢。
- timesteps (list[int], 可选) — 用于支持 timesteps 参数的调度器去噪过程的自定义时间步，该参数在其 set_timesteps 方法中使用。如果未定义，则使用传入 num_inference_steps 时的默认行为。必须按降序排列。
- guidance_scale (float，可选，默认为 4.0) — Classifier-Free Diffusion Guidance 中定义的引导尺度。guidance_scale 定义为 Imagen Paper 中公式 2 的 w。通过设置 guidance_scale > 1 来启用引导尺度。较高的引导尺度会促使生成与文本 prompt 紧密关联的图像，通常以牺牲图像质量为代价。
- num_images_per_prompt (int，可选，默认为 1) — 每个提示词生成的图像数量。
- generator (torch.Generator 或 list[torch.Generator]，可选) — 一个或一组 torch 生成器，用于使生成过程具有确定性。
- latents (torch.Tensor, optional) — 预先生成的含噪潜变量，从高斯分布中采样得到，用作图像生成的输入。可用于以不同的提示词微调同一生成过程。如果未提供，将通过使用提供的随机 generator 进行采样来生成一个潜变量张量。
- prompt_embeds (torch.FloatTensor, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入，例如 提示词加权。如果未提供，文本嵌入将从 prompt 输入参数生成。
- negative_prompt_embeds (torch.FloatTensor，可选) — 预生成的负向文本嵌入。可用于轻松调整文本输入，例如提示词加权。如果未提供且 guidance_scale > 1，则将从空字符串生成负向嵌入。
- prompt_attention_mask (torch.BoolTensor，可选) — prompt_embeds 的预生成注意力掩码。如果未提供，则会根据 prompt 输入参数生成注意力掩码。
- negative_prompt_attention_mask (torch.BoolTensor, 可选) — negative_prompt_embeds 的预生成注意力掩码。如果未提供且 guidance_scale > 1，则将从空字符串生成注意力掩码。
- output_type (str，可选，默认为 “pil”) — 生成图像的输出格式。可选值为 PIL：PIL.Image.Image 或 np.array。
- tokenizer_max_length (int, 可选) —
覆盖对提示词进行分词时使用的最大 token 数。如果未设置，则默认为分词器自身的
model_max_length。 - skip_text_cleaning (bool，可选，默认为 False) — 如果为 True，则仅使用轻量级提示清理（修复编码 + 反转义 HTML），而不是完整的 DeepFloyd 清理流程。
- return_dict (bool，可选，默认为 True) — 是否返回 [~pipelines.prx.PRXPipelineOutput] 而不是普通元组。
- use_resolution_binning (bool，可选，默认为 True) — 如果设置为 True，请求的高度和宽度会先通过预定义的宽高比分箱映射到最接近的分辨率。在生成的潜变量解码为图像后，它们会被重新调整回请求的分辨率。适用于以最佳分辨率生成非正方形图像。
- callback_on_step_end (Callable，可选) — 在推理过程中，每个去噪步骤结束时调用的函数。该函数以以下参数被调用：callback_on_step_end(self, step, timestep, callback_kwargs)。 callback_kwargs 将包含由 callback_on_step_end_tensor_inputs 指定的所有张量列表。
- callback_on_step_end_tensor_inputs (list，optional) — callback_on_step_end 函数的张量输入列表。列表中指定的张量将作为 callback_kwargs 参数传入。您只能包含在 ._callback_tensor_inputs 属性中列出的张量。
返回值
[~pipelines.prx.PRXPipelineOutput] 或 *tuple`
[~pipelines.prx.PRXPipelineOutput] 如果 return_dict 为 True，否则为 *tuple。当返回元组时，第一个元素是一个包含生成图像的列表。
调用生成 pipeline 时调用的函数。
示例：
>>> import torch
>>> from diffusers import PRXPipeline
>>> # Load pipeline with from_pretrained
>>> pipe = PRXPipeline.from_pretrained("Photoroom/prx-512-t2i-sft")
>>> pipe.to("cuda")
>>> prompt = "A digital painting of a rusty, vintage tram on a sandy beach"
>>> image = pipe(prompt, num_inference_steps=28, guidance_scale=5.0).images[0]
>>> image.save("prx_output.png")check_inputs
< source >( prompt: str | list[str]height: intwidth: intguidance_scale: floatcallback_on_step_end_tensor_inputs: list[str] | None = Noneprompt_embeds: typing.Optional[torch.FloatTensor] = Nonenegative_prompt_embeds: typing.Optional[torch.FloatTensor] = None )
检查所有输入是否处于正确的格式。
encode_prompt
< source >( prompt: str | list[str]device: typing.Optional[torch.device] = Nonedo_classifier_free_guidance: bool = Truenegative_prompt: str = ''num_images_per_prompt: int = 1prompt_embeds: typing.Optional[torch.FloatTensor] = Nonenegative_prompt_embeds: typing.Optional[torch.FloatTensor] = Noneprompt_attention_mask: typing.Optional[torch.BoolTensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.BoolTensor] = Nonetokenizer_max_length: int | None = Noneskip_text_cleaning: bool = False )
使用标准文本编码器和分词器对文本提示进行编码，或使用预计算的嵌入。
根据加载的 VAE 和配置确定默认分辨率。
prepare_latents
< source >( batch_size: intnum_channels_latents: intheight: intwidth: intdtype: dtypedevice: devicegenerator: typing.Optional[torch.Generator] = Nonelatents: typing.Optional[torch.Tensor] = None )
为扩散过程准备初始潜变量。
PRXPipelineOutput
class diffusers.pipelines.prx.PRXPipelineOutput
< source >( images: list[PIL.Image.Image] | numpy.ndarray )
PRX 流水线的输出类。