#
# 根据 Apache 许可证 2.0 版本（“许可证”）授权
# 您只能在遵守许可证的情况下使用本文件。
# 您可以在以下位置获取许可证副本
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 除非适用法律要求或双方书面同意，软件
# 根据许可证分发，本软件按“原样”
# 不附带任何明示或暗示的保证或条件。
# 请参阅许可证中关于权限的具体语言
# 许可证下的限制。-->

# PRX

PRX 使用简化的 MMDIT 架构从文本生成高质量图像，在该架构中，文本 token 不会通过 transformer 块进行更新。它采用带有离散调度的 flow matching 以实现高效采样，并使用 Google 的 T5Gemma-2B-2B-UL2 模型进行多语言文本编码。该约 1.3B 参数的 transformer 在不牺牲质量的前提下实现了快速推理。您可以选择 Flux VAE（8 倍压缩，16 个潜在通道）以平衡质量和速度，或选择 DC-AE（32 倍压缩，32 个潜在通道）以实现潜在压缩和更快的处理。

## 可用模型

PRX 提供多种具有不同 VAE 配置的变体，每种变体均针对特定分辨率进行了优化。基础模型在处理详细提示词方面表现出色，能够捕捉复杂的构图和细微的细节。在 [Alchemist 数据集](https://huggingface.co/datasets/yandex/alchemist) 上经过微调的模型提升了美学质量，尤其是在使用较简单的提示词时。

| 模型 | 分辨率 | 微调 | 蒸馏 | 描述 | 建议的提示 | 建议参数 | 推荐的 dtype |
|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|
| [`Photoroom/prx-256-t2i`](https://huggingface.co/Photoroom/prx-256-t2i)| 256 | 否 | 否 | 使用 Flux VAE 在 256 分辨率下预训练的基础模型|在自然语言中提供详细的提示效果最佳|28 步，cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-256-t2i-sft`](https://huggingface.co/Photoroom/prx-256-t2i-sft)| 512 | 是的 | 否 | 使用 Flux VAE 在 [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) 数据集上进行微调 | 能够处理细节较少的提示|28 步，cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i`](https://huggingface.co/Photoroom/prx-512-t2i)| 512 | 否 | 否 | 使用 Flux VAE 在 512 分辨率下预训练的基础模型 |在自然语言中提供详细的提示效果最佳|28 步，cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft)| 512 | 是的 | 否 | 使用 Flux VAE 在 [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) 数据集上进行微调 | 能够处理自然语言中细节较少的提示|28 步，cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-sft-distilled)| 512 | 是的 | 是的 | 来自 [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft) 的 8 步蒸馏模型 | 能够处理自然语言中细节较少的提示|8 步，cfg=1.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae)| 512 | 否 | 否 | 使用 [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) 在 512 分辨率下预训练的基础模型|在自然语言中提供详细的提示效果最佳|28 步，cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft)| 512 | 是的 | 否 | 使用 [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) 在 [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) 数据集上进行微调 | 能够处理自然语言中细节较少的提示|28 步，cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled)| 512 | 是 | 是 | 从 [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled) 蒸馏而来的 8 步模型 | 能够处理自然语言中细节较少的提示词 |8 步，cfg=1.0| `torch.bfloat16` |s

有关更多信息，请参阅[此](https://huggingface.co/collections/Photoroom/prx-models-68e66254c202ebfab99ad38e)集合。

## 加载 pipeline

使用 [from_pretrained()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained) 加载 pipeline。

```py
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")
```

### 手动加载组件

单独加载各个组件，以自定义流水线，例如使用量化模型。

```py
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"
```

## 内存优化

对于内存受限的环境：

```py
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[[diffusers.PRXPipeline]]

#### diffusers.PRXPipeline[[diffusers.PRXPipeline]]

```python
diffusers.PRXPipeline(transformer: PRXTransformer2DModel, scheduler: FlowMatchEulerDiscreteScheduler, text_encoder: T5GemmaEncoder, tokenizer: transformers.models.t5.tokenization_t5.T5Tokenizer | transformers.models.gemma.tokenization_gemma.GemmaTokenizer | transformers.models.auto.tokenization_auto.AutoTokenizer, vae: diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL | diffusers.models.autoencoders.autoencoder_dc.AutoencoderDC | None = None, default_sample_size: int | None = 512)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/prx/pipeline_prx.py#L257)

**参数：**

transformer (`PRXTransformer2DModel`)：用于对编码后的图像潜变量进行去噪的 PRX transformer 模型。

scheduler ([FlowMatchEulerDiscreteScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler))：用于与 `transformer` 配合，对编码后的图像潜变量进行去噪的调度器。

text_encoder (`T5GemmaEncoder`)：用于编码提示词的文本编码器模型。

tokenizer （`T5TokenizerFast` 或 `GemmaTokenizerFast`）：用于文本编码器的分词器。

vae ([AutoencoderKL](/docs/diffusers/v0.43.0/zh/api/models/autoencoderkl#diffusers.AutoencoderKL) 或 [AutoencoderDC](/docs/diffusers/v0.43.0/zh/api/models/autoencoder_dc#diffusers.AutoencoderDC))：变分自编码器（VAE）模型，用于将图像编码为潜在表示并从中解码。同时支持 AutoencoderKL（8 倍压缩）和 AutoencoderDC（32 倍压缩）。

使用 PRX Transformer 进行文本到图像生成的 Pipeline。

该模型继承自 [DiffusionPipeline](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline)。请查阅父类文档，了解该库为所有 pipeline 实现的通用方法（例如下载或保存、在特定设备上运行等）。

#### __调用__[[diffusers.PRXPipeline.__call__]]

```python
__call__(prompt: str | list[str] = None, negative_prompt: str = '', height: int | None = None, width: int | None = None, num_inference_steps: int = 28, timesteps: list[int] = None, guidance_scale: float = 4.0, num_images_per_prompt: int | None = 1, generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = None, latents: typing.Optional[torch.Tensor] = None, prompt_embeds: typing.Optional[torch.FloatTensor] = None, negative_prompt_embeds: typing.Optional[torch.FloatTensor] = None, prompt_attention_mask: typing.Optional[torch.BoolTensor] = None, negative_prompt_attention_mask: typing.Optional[torch.BoolTensor] = None, output_type: str | None = 'pil', return_dict: bool = True, use_resolution_binning: bool = True, callback_on_step_end: typing.Optional[typing.Callable[[int, int], NoneType]] = None, callback_on_step_end_tensor_inputs: list[str] = ['latents'], tokenizer_max_length: int | None = None, skip_text_cleaning: bool = False)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/prx/pipeline_prx.py#L554)

**参数：**

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*，*optional*，默认为 self.transformer.config.sample_size * self.vae_scale_factor)：生成图像的像素宽度。

num_inference_steps (*int*, *optional*，默认为 28)：去噪步数。更多的去噪步数通常能生成更高质量的图像，但代价是推理速度变慢。

timesteps (*list[int]*, *optional*) : 用于支持在 *set_timesteps* 方法中接受 *timesteps* 参数的调度器的去噪过程的自定义时间步。如果未定义，将使用传入 *num_inference_steps* 时的默认行为。必须按降序排列。

guidance_scale (*float*, *optional*，默认为 4.0)：如 [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598) 中定义的引导尺度。*guidance_scale* 定义为 [Imagen Paper](https://huggingface.co/papers/2205.11487) 中公式 2 的 *w*。通过设置 *guidance_scale > 1* 来启用引导尺度。较高的引导尺度会促使生成与文本 *prompt* 紧密关联的图像，通常以牺牲图像质量为代价。

num_images_per_prompt (*int*，*optional*，默认为 1)：每个提示词生成的图像数量。

generator (*torch.Generator* 或 *list[torch.Generator]*，*可选*)：一个或一组 [torch 生成器](https://pytorch.org/docs/stable/generated/torch.Generator.html)，用于使生成过程具有确定性。

latents (*torch.Tensor*, *可选*) : 预先生成的噪声潜变量，从高斯分布中采样，用作图像生成的输入。可用于使用不同的提示词对同一生成过程进行微调。如果未提供，将通过使用提供的随机 *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*，*optional*，默认为 *"pil"*)：生成图像的输出格式。可选 [PIL](https://pillow.readthedocs.io/en/stable/)：*PIL.Image.Image* 或 *np.array*。

tokenizer_max_length (*int*, *optional*) : 覆盖对提示词进行分词时使用的最大 token 数。如果未设置，则默认为分词器自身的 `model_max_length`。

skip_text_cleaning (*bool*，*可选*，默认为 *False*)：如果为 *True*，则仅使用轻量级提示词清理（修复编码 + 反转义 HTML），而不是完整的 DeepFloyd 清理流程。

return_dict (*bool*，*optional*，默认为 *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 时调用的函数。

示例：
```py
>>> 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[[diffusers.PRXPipeline.check_inputs]]

```python
check_inputs(prompt: str | list[str], height: int, width: int, guidance_scale: float, callback_on_step_end_tensor_inputs: list[str] | None = None, prompt_embeds: typing.Optional[torch.FloatTensor] = None, negative_prompt_embeds: typing.Optional[torch.FloatTensor] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/prx/pipeline_prx.py#L500)

检查所有输入是否处于正确的格式。

#### encode_prompt[[diffusers.PRXPipeline.encode_prompt]]

```python
encode_prompt(prompt: str | list[str], device: typing.Optional[torch.device] = None, do_classifier_free_guidance: bool = True, negative_prompt: str = '', num_images_per_prompt: int = 1, prompt_embeds: typing.Optional[torch.FloatTensor] = None, negative_prompt_embeds: typing.Optional[torch.FloatTensor] = None, prompt_attention_mask: typing.Optional[torch.BoolTensor] = None, negative_prompt_attention_mask: typing.Optional[torch.BoolTensor] = None, tokenizer_max_length: int | None = None, skip_text_cleaning: bool = False)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/prx/pipeline_prx.py#L376)

使用标准文本编码器和分词器对文本提示进行编码，或使用预计算的嵌入。

#### get_default_resolution[[diffusers.PRXPipeline.get_default_resolution]]

```python
get_default_resolution()
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/prx/pipeline_prx.py#L340)

**返回值：** `int`

用于生成的默认样本尺寸（高度/宽度）。

根据加载的 VAE 和配置确定默认分辨率。

#### prepare_latents[[diffusers.PRXPipeline.prepare_latents]]

```python
prepare_latents(batch_size: int, num_channels_latents: int, height: int, width: int, dtype: dtype, device: device, generator: typing.Optional[torch.Generator] = None, latents: typing.Optional[torch.Tensor] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/prx/pipeline_prx.py#L352)

为扩散过程准备初始潜变量。

## PRXPipelineOutput[[diffusers.pipelines.prx.PRXPipelineOutput]]

#### diffusers.pipelines.prx.PRXPipelineOutput[[diffusers.pipelines.prx.PRXPipelineOutput]]

```python
diffusers.pipelines.prx.PRXPipelineOutput(images: list[PIL.Image.Image] | numpy.ndarray)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/prx/pipeline_output.py#L24)

**参数：**

images (`list[PIL.Image.Image]` 或 `np.ndarray`)：长度为 `batch_size` 的去噪 PIL 图像列表，或形状为 `(batch_size, height, width, num_channels)` 的 numpy 数组。PIL 图像或 numpy 数组表示扩散流水线生成的去噪图像。

PRX 流水线的输出类。