#
# Licenciado sob a Apache License, Versão 2.0 (a "Licença");
# você não pode usar este arquivo, exceto em conformidade com a Licença.
# Você pode obter uma cópia da Licença em
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# A menos que exigido pela lei aplicável ou acordado por escrito, o software
# distribuído sob a Licença é distribuído em uma base "COMO ESTÁ",
# SEM GARANTIAS OU CONDIÇÕES DE QUALQUER TIPO, expressas ou implícitas.
# Consulte a Licença para o idioma específico que rege as permissões e
# limitações sob a Licença. -->

# PRX

O PRX gera imagens de alta qualidade a partir de texto usando uma arquitetura MMDIT simplificada, na qual os tokens de texto não são atualizados através dos blocos do transformer. Ele emprega flow matching com agendamento discreto para amostragem eficiente e utiliza o modelo T5Gemma-2B-2B-UL2 da Google para codificação de texto multilíngue. O transformer de ~1.3B parâmetros oferece inferência rápida sem sacrificar a qualidade. Você pode escolher entre o Flux VAE (compressão 8x, 16 canais latentes) para um equilíbrio entre qualidade e velocidade, ou o DC-AE (compressão 32x, 32 canais latentes) para compressão latente e processamento mais rápido.

## Modelos disponíveis

O PRX oferece múltiplas variantes com diferentes configurações de VAE, cada uma otimizada para resoluções específicas. Os modelos base se destacam com prompts detalhados, capturando composições complexas e detalhes sutis. Modelos ajustados e treinados no [conjunto de dados Alchemist](https://huggingface.co/datasets/yandex/alchemist) melhoram a qualidade estética, especialmente com prompts mais simples.

| Model | Resolução | Ajustado | Destilado | Descrição | Prompts sugeridos | Parâmetros sugeridos | dtype recomendado |
|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|
| [`Photoroom/prx-256-t2i`](https://huggingface.co/Photoroom/prx-256-t2i)| 256 | Não | Não | Modelo base pré-treinado em 256 com Flux VAE|Funciona melhor com prompts detalhados em linguagem natural|28 etapas, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-256-t2i-sft`](https://huggingface.co/Photoroom/prx-256-t2i-sft)| 512 | Sim | Não | Ajuste fino realizado no conjunto de dados [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) com Flux VAE | Consegue lidar com prompts menos detalhados|28 etapas, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i`](https://huggingface.co/Photoroom/prx-512-t2i)| 512 | Não | Não | Modelo base pré-treinado em 512 com Flux VAE |Funciona melhor com prompts detalhados em linguagem natural|28 etapas, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft)| 512 | Sim | Não | Ajuste fino realizado no conjunto de dados [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) com Flux VAE | Consegue lidar com prompts menos detalhados em linguagem natural|28 etapas, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-sft-distilled)| 512 | Sim | Sim | Modelo destilado de 8 etapas de [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft) | Consegue lidar com prompts menos detalhados em linguagem natural|8 etapas, cfg=1.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae)| 512 | Não | Não | Modelo base pré-treinado em 512 com [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)|Funciona melhor com prompts detalhados em linguagem natural|28 etapas, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft)| 512 | Sim | Não | Ajustado no conjunto de dados [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) com [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) | Consegue lidar com prompts menos detalhados em linguagem natural|28 etapas, 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 | Sim | Sim | Modelo destilado de 8 etapas a partir de [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled) | Pode lidar com prompts menos detalhados em linguagem natural|8 etapas, cfg=1.0| `torch.bfloat16` |s

Consulte esta [coleção](https://huggingface.co/collections/Photoroom/prx-models-68e66254c202ebfab99ad38e) para mais informações.

## Carregando o pipeline

Load the pipeline with [from_pretrained()](/docs/diffusers/v0.43.0/pt/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")
```

### Carregamento Manual de Componentes

Carregue os componentes individualmente para personalizar o pipeline, por exemplo, para usar modelos quantizados.

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

## Otimização de Memória

Para ambientes com restrições de memória:

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

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

**Parâmetros:**

transformer (`PRXTransformer2DModel`) : O modelo transformer PRX para desruidar os latentes da imagem codificada.

scheduler ([FlowMatchEulerDiscreteScheduler](/docs/diffusers/v0.43.0/pt/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler)) : Um scheduler a ser usado em combinação com `transformer` para desruidar os latentes da imagem codificada.

text_encoder (`T5GemmaEncoder`) : Modelo de codificador de texto para codificar prompts.

tokenizer ([`T5TokenizerFast` ou `GemmaTokenizerFast`]) : Tokenizer para o codificador de texto.

vae ([AutoencoderKL](/docs/diffusers/v0.43.0/pt/api/models/autoencoderkl#diffusers.AutoencoderKL) ou [AutoencoderDC](/docs/diffusers/v0.43.0/pt/api/models/autoencoder_dc#diffusers.AutoencoderDC)) : Modelo de Variational Auto-Encoder (VAE) para codificar e decodificar imagens para e a partir de representações latentes. Suporta tanto AutoencoderKL (compressão 8x) quanto AutoencoderDC (compressão 32x).

Pipeline para geração de texto para imagem usando PRX Transformer.

Este modelo herda de [DiffusionPipeline](/docs/diffusers/v0.43.0/pt/api/pipelines/overview#diffusers.DiffusionPipeline). Consulte a documentação da superclasse para os métodos genéricos que a biblioteca implementa para todos os pipelines (como download ou salvamento, execução em um dispositivo específico, etc.)

#### __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)
```

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

**Parâmetros:**

prompt (*str* ou *list[str]*, *opcional*) : O prompt ou os prompts para orientar a geração da imagem. Se não definido, é necessário passar *prompt_embeds* em vez disso.

negative_prompt (*str*, *opcional*, padrão *""*) : O prompt ou os prompts que não devem orientar a geração da imagem. Ignorado quando não se usa orientação (ou seja, ignorado se *guidance_scale* for menor que *1*).

height (*int*, *optional*, padrão é self.transformer.config.sample_size * self.vae_scale_factor) : A altura em pixels da imagem gerada.

width (*int*, *optional*, padrão é self.transformer.config.sample_size * self.vae_scale_factor) : A largura em pixels da imagem gerada.

num_inference_steps (*int*, *optional*, padrão 28) : O número de etapas de desruidamento. Mais etapas de desruidamento geralmente resultam em uma imagem de maior qualidade, em detrimento de uma inferência mais lenta.

timesteps (*list[int]*, *optional*) : Passos de tempo personalizados para usar no processo de remoção de ruído com agendadores que suportam um argumento *timesteps* em seu método *set_timesteps*. Se não definido, o comportamento padrão quando *num_inference_steps* é passado será usado. Deve estar em ordem decrescente.

guidance_scale (*float*, *optional*, padrão 4.0) : Escala de orientação definida em [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598). *guidance_scale* é definido como *w* da equação 2. do [Imagen Paper](https://huggingface.co/papers/2205.11487). A escala de orientação é ativada ao definir *guidance_scale > 1*. Uma escala de orientação maior incentiva a geração de imagens que estão intimamente ligadas ao *prompt* de texto, geralmente em detrimento de uma menor qualidade de imagem.

num_images_per_prompt (*int*, *optional*, padrão 1) : O número de imagens a serem geradas por prompt.

generator (*torch.Generator* ou *list[torch.Generator]*, *opcional*) : Um ou uma lista de [geradores do torch](https://pytorch.org/docs/stable/generated/torch.Generator.html) para tornar a geração determinística.

latents (*torch.Tensor*, *opcional*) : Latentes ruidosos pré-gerados, amostrados de uma distribuição gaussiana, a serem usados como entradas para a geração de imagens. Pode ser usado para ajustar a mesma geração com diferentes prompts. Se não for fornecido, um tensor de latentes será gerado por amostragem usando o *generator* aleatório fornecido.

prompt_embeds (*torch.FloatTensor*, *opcional*) : Embebimentos de texto pré-gerados. Pode ser usado para ajustar facilmente as entradas de texto, *por exemplo*, ponderação de prompts. Se não for fornecido, os embebimentos de texto serão gerados a partir do argumento de entrada *prompt*.

negative_prompt_embeds (*torch.FloatTensor*, *opcional*) : Embeddings de texto negativos pré-gerados. Pode ser usado para ajustar facilmente as entradas de texto, *por exemplo*, ponderação de prompts. Se não for fornecido e *guidance_scale > 1*, os embeddings negativos serão gerados a partir de uma string vazia.

prompt_attention_mask (*torch.BoolTensor*, *opcional*) : Máscara de atenção pré-gerada para *prompt_embeds*. Se não for fornecida, a máscara de atenção será gerada a partir do argumento de entrada *prompt*.

negative_prompt_attention_mask (*torch.BoolTensor*, *opcional*) : Máscara de atenção pré-gerada para *negative_prompt_embeds*. Se não for fornecida e *guidance_scale > 1*, a máscara de atenção será gerada a partir de uma string vazia.

output_type (*str*, *opcional*, padrão *"pil"*) : O formato de saída da imagem gerada. Escolha entre [PIL](https://pillow.readthedocs.io/en/stable/): *PIL.Image.Image* ou *np.array*.

tokenizer_max_length (*int*, *opcional*) : Substitui o número máximo de tokens usados ao tokenizar o prompt. Padrão é o próprio `model_max_length` do tokenizer, quando não definido.

skip_text_cleaning (*bool*, *opcional*, padrão *False*) : Se *True*, usa apenas a limpeza leve do prompt (correção de codificação + desescape de HTML) em vez do pipeline completo de limpeza do DeepFloyd.

return_dict (*bool*, *opcional*, padrão *True*) : Se deve ou não retornar um [*~pipelines.prx.PRXPipelineOutput*] em vez de uma tupla simples.

use_resolution_binning (*bool*, *opcional*, padrão *True*) : Se definido como *True*, a altura e a largura solicitadas são primeiro mapeadas para as resoluções mais próximas usando bins de proporção de aspecto predefinidos. Após os latentes gerados serem decodificados em imagens, eles são redimensionados de volta para a resolução solicitada. Útil para gerar imagens não quadradas em resoluções ótimas.

callback_on_step_end (*Callable*, *opcional*) : Uma função que é chamada no final de cada etapa de desruidamento durante a inferência. A função é chamada com os seguintes argumentos: *callback_on_step_end(self, step, timestep, callback_kwargs)*. *callback_kwargs* incluirá uma lista de todos os tensores, conforme especificado por *callback_on_step_end_tensor_inputs*.

callback_on_step_end_tensor_inputs (*list*, *optional*) : A lista de entradas de tensor para a função *callback_on_step_end*. Os tensores especificados na lista serão passados como argumento *callback_kwargs*. Você só poderá incluir tensores que estão listados no atributo *._callback_tensor_inputs*.

**Retorna:** [*~pipelines.prx.PRXPipelineOutput*] ou *tuple`

[*~pipelines.prx.PRXPipelineOutput*] se *return_dict* for True, caso contrário, uma *tuple. Ao retornar uma tuple, o primeiro elemento é uma lista com as imagens geradas.

Função invocada ao chamar o pipeline para geração.

Exemplos:
```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)
```

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

Verifique se todas as entradas estão no formato correto.

#### 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)
```

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

Codifique o prompt de texto usando um codificador e um tokenizador de texto padrão, ou use embeddings pré-computados.

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

```python
get_default_resolution()
```

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

**Retorna:** `int`

O tamanho padrão da amostra (altura/largura) a ser usado para a geração.

Determine a resolução padrão com base no VAE carregado e na configuração.

#### 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)
```

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

Prepare os latentes iniciais para o processo de difusão.

## 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)
```

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

**Parâmetros:**

images (`list[PIL.Image.Image]` ou `np.ndarray`) : lista de imagens PIL desruidas de comprimento `batch_size` ou array numpy de forma `(batch_size, height, width, num_channels)`. As imagens PIL ou o array numpy apresentam as imagens desruidas do pipeline de difusão.

Classe de saída para pipelines PRX.