#
# Лицензия Apache License, Version 2.0 («Лицензия»);
# вы можете использовать этот файл только в соответствии с Лицензией.
# Вы можете получить копию Лицензии по адресу
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Если это не требуется применимым законодательством или не оговорено письменно, программное
# распространяется по Лицензии на условиях «КАК ЕСТЬ»,
# БЕЗ КАКИХ-ЛИБО ГАРАНТИЙ ИЛИ УСЛОВИЙ, прямо указанных или подразумеваемых.
# См. Лицензию для получения конкретных условий, регулирующих разрешения и
# ограничениями Лицензии. -->

# PRX

PRX генерирует изображения высокого качества по тексту, используя упрощённую архитектуру MMDIT, в которой токены текста не обновляются через блоки трансформера. Для эффективного сэмплирования применяется flow matching с дискретным планированием, а для многоязычного кодирования текста используется модель T5Gemma-2B-2B-UL2 от Google. Трансформер с ~1,3 млрд параметров обеспечивает быстрое инференсирование без потери качества. Вы можете выбрать между Flux VAE (сжатие 8x, 16 латентных каналов) для баланса качества и скорости или DC-AE (сжатие 32x, 32 латентных канала) для сжатия латентного пространства и более быстрой обработки.

## Доступные модели

PRX предлагает несколько вариантов с различными конфигурациями VAE, каждый из которых оптимизирован для конкретных разрешений. Базовые модели лучше всего работают с детализированными промптами, позволяя улавливать сложные композиции и тонкие детали. Модели, дообученные на [наборе данных Alchemist](https://huggingface.co/datasets/yandex/alchemist), повышают эстетическое качество, особенно при использовании более простых промптов.

| Модель | Разрешение | Доработанные | Дистиллированный | Описание | Рекомендуемые промпты | Рекомендуемые параметры | Рекомендуемый dtype |
|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|
| [`Photoroom/prx-256-t2i`](https://huggingface.co/Photoroom/prx-256-t2i)| 256 | Нет | Нет | Базовая модель, предварительно обученная при 256 с VAE Flux|Лучше всего работает с подробными запросами на естественном языке|28 шагов, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-256-t2i-sft`](https://huggingface.co/Photoroom/prx-256-t2i-sft)| 512 | Да | Нет | Доработан на наборе данных [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) с использованием Flux VAE | Способен обрабатывать менее детализированные запросы|28 шагов, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i`](https://huggingface.co/Photoroom/prx-512-t2i)| 512 | Нет | Нет | Базовая модель, предварительно обученная при разрешении 512 с VAE Flux |Лучше всего работает с подробными запросами на естественном языке|28 шагов, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft)| 512 | Да | Нет | Доработан на наборе данных [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) с использованием Flux VAE | Способен обрабатывать менее детализированные запросы на естественном языке|28 шагов, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-sft-distilled)| 512 | Да | Да | 8-шаговая дистиллированная модель из [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft) | Способен обрабатывать менее детализированные запросы на естественном языке|8 шагов, cfg=1.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae)| 512 | Нет | Нет | Базовая модель, предварительно обученная при разрешении 512 с использованием [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)|Лучше всего работает с подробными запросами на естественном языке|28 шагов, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft)| 512 | Да | Нет | Доработан на наборе данных [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) с использованием [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) | Способен обрабатывать менее детализированные запросы на естественном языке|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 | Да | Да | 8-шаговая дистиллированная модель из [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled) | Может обрабатывать менее детализированные промпты на естественном языке|8 шагов, cfg=1.0| `torch.bfloat16` |s

См. [это](https://huggingface.co/collections/Photoroom/prx-models-68e66254c202ebfab99ad38e) собрание для получения дополнительной информации.

## Загрузка конвейера

Загрузите конвейер с помощью [from_pretrained()](/docs/diffusers/v0.43.0/ru/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained).

```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/ru/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler)) : планировщик, используемый в связке с `transformer` для удаления шума из закодированных латентов изображения.

text_encoder (`T5GemmaEncoder`) : Текстовая модель-энкодер для кодирования запросов.

tokenizer ([`T5TokenizerFast` или `GemmaTokenizerFast`]) : Токенизатор для текстового энкодера.

vae ([AutoencoderKL](/docs/diffusers/v0.43.0/ru/api/models/autoencoderkl#diffusers.AutoencoderKL) или [AutoencoderDC](/docs/diffusers/v0.43.0/ru/api/models/autoencoder_dc#diffusers.AutoencoderDC)) : Модель вариационного автоэнкодера (VAE) для кодирования изображений в латентные представления и декодирования из них. Поддерживает как AutoencoderKL (сжатие 8x), так и AutoencoderDC (сжатие 32x).

Пайплайн для генерации изображений по тексту с использованием PRX Transformer.

Эта модель наследуется от [DiffusionPipeline](/docs/diffusers/v0.43.0/ru/api/pipelines/overview#diffusers.DiffusionPipeline). См. документацию родительского класса для общих методов, которые библиотека реализует для всех пайплайнов (например, загрузка или сохранение, запуск на конкретном устройстве и т. д.)

#### __call__[[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*, *optional*, по умолчанию *""*) : Промпт или промпты, которые не должны направлять генерацию изображения. Игнорируется, если не используется направляющая (то есть игнорируется, если *guidance_scale* меньше *1*).

height (*int*, *optional*, по умолчанию 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*) : Пользовательские шаги для процесса денойзинга с планировщиками, поддерживающими аргумент *timesteps* в их методе *set_timesteps*. Если не определено, будет использоваться поведение по умолчанию при передаче *num_inference_steps*. Должны быть в порядке убывания.

guidance_scale (*float*, *optional*, по умолчанию 4.0) : Масштаб руководства, определенный в [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598). *guidance_scale* определяется как *w* из уравнения 2 в [Imagen Paper](https://huggingface.co/papers/2205.11487). Масштаб руководства активируется при установке *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*, *optional*) : Предварительно сгенерированные зашумлённые латенты, выбранные из гауссова распределения, которые используются в качестве входных данных для генерации изображений. Можно использовать, чтобы скорректировать одну и ту же генерацию с разными промптами. Если не предоставлены, тензор латентов будет сгенерирован путём выборки с использованием переданного случайного *generator*.

prompt_embeds (*torch.FloatTensor*, *optional*) : Предварительно сгенерированные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, *например* для взвешивания промпта. Если не предоставлены, текстовые эмбеддинги будут сгенерированы из входного аргумента *prompt*.

negative_prompt_embeds (*torch.FloatTensor*, *optional*) : Предварительно сгенерированные отрицательные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, *например* для взвешивания промптов. Если не предоставлены и *guidance_scale > 1*, отрицательные эмбеддинги будут сгенерированы из пустой строки.

prompt_attention_mask (*torch.BoolTensor*, *optional*) : Предварительно сгенерированная маска внимания для *prompt_embeds*. Если не предоставлена, маска внимания будет сгенерирована из аргумента *prompt*.

negative_prompt_attention_mask (*torch.BoolTensor*, *optional*) : Предварительно сгенерированная маска внимания для *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*) : Переопределяет максимальное количество токенов, используемых при токенизации промпта. По умолчанию используется `model_max_length` самого токенизатора, если значение не задано.

skip_text_cleaning (*bool*, *optional*, по умолчанию *False*) : Если *True*, используется только лёгкая очистка промпта (исправление кодировки + снятие HTML-экранирования) вместо полного конвейера очистки DeepFloyd.

return_dict (*bool*, *optional*, по умолчанию *True*) : Возвращать ли [*~pipelines.prx.PRXPipelineOutput*] вместо обычного кортежа.

use_resolution_binning (*bool*, *optional*, по умолчанию *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, иначе — *кортеж. При возврате кортежа первым элементом является список с сгенерированными изображениями.

Функция, вызываемая при обращении к конвейеру для генерации.

Примеры:
```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`) : список денойзированных изображений PIL длиной `batch_size` или массив numpy с формой `(batch_size, height, width, num_channels)`. Изображения PIL или массив numpy представляют денойзированные изображения диффузионного конвейера.

Класс вывода для конвейеров PRX.