#
# Bajo licencia Apache License, versión 2.0 (la "Licencia");
# no puede usar este archivo excepto en cumplimiento de la Licencia.
# Puede obtener una copia de la Licencia en
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# A menos que la ley aplicable lo exija o se haya acordado por escrito, el software
# distribuido bajo la Licencia se distribuye en una base "TAL CUAL",
# SIN GARANTÍAS NI CONDICIONES DE NINGÚN TIPO, ya sea expresa o implícita.
# Consulte la Licencia para el lenguaje específico que rige los permisos y
# limitaciones de la Licencia. -->

# PRX

PRX genera imágenes de alta calidad a partir de texto utilizando una arquitectura MMDIT simplificada en la que los tokens de texto no se actualizan a través de los bloques transformer. Emplea flow matching con programación discreta para una muestreo eficiente y utiliza el modelo T5Gemma-2B-2B-UL2 de Google para la codificación de texto multilingüe. El transformer de ~1.3B parámetros ofrece una inferencia rápida sin sacrificar la calidad. Puedes elegir entre Flux VAE (compresión 8x, 16 canales latentes) para un equilibrio entre calidad y velocidad, o DC-AE (compresión 32x, 32 canales latentes) para una mayor compresión latente y un procesamiento más rápido.

## Modelos disponibles

PRX ofrece múltiples variantes con diferentes configuraciones de VAE, cada una optimizada para resoluciones específicas. Los modelos base destacan con prompts detallados, capturando composiciones complejas y sutiles detalles. Los modelos afinados entrenados en el [conjunto de datos Alchemist](https://huggingface.co/datasets/yandex/alchemist) mejoran la calidad estética, especialmente con prompts más simples.

| Modelo | Resolución | Ajustado | Destilado | Descripción | Prompts sugeridos | Parámetros sugeridos | dtype recomendado |
|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|
| [`Photoroom/prx-256-t2i`](https://huggingface.co/Photoroom/prx-256-t2i)| 256 | No | No | Modelo base preentrenado en 256 con Flux VAE|Funciona mejor con prompts detallados en lenguaje natural|28 pasos, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-256-t2i-sft`](https://huggingface.co/Photoroom/prx-256-t2i-sft)| 512 | Sí | No | Ajustado fino en el conjunto de datos [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) con Flux VAE | Puede manejar prompts menos detallados|28 pasos, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i`](https://huggingface.co/Photoroom/prx-512-t2i)| 512 | No | No | Modelo base preentrenado en 512 con Flux VAE |Funciona mejor con prompts detallados en lenguaje natural|28 pasos, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft)| 512 | Sí | No | Ajustado fino en el conjunto de datos [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) con Flux VAE | Puede manejar prompts menos detallados en lenguaje natural|28 pasos, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-sft-distilled)| 512 | Sí | Sí | Modelo destilado de 8 pasos de [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft) | Puede manejar prompts menos detallados en lenguaje natural|8 pasos, cfg=1.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae)| 512 | No | No | Modelo base preentrenado a 512 con [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)|Funciona mejor con prompts detallados en lenguaje natural|28 pasos, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft)| 512 | Sí | No | Ajustado finamente en el [conjunto de datos Alchemist](https://huggingface.co/datasets/yandex/alchemist) con [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) | Puede manejar prompts menos detallados en lenguaje natural|28 pasos, 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 | Sí | Sí | Modelo destilado de 8 pasos a partir de [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled) | Puede manejar prompts menos detallados en lenguaje natural|8 pasos, cfg=1.0| `torch.bfloat16` |s

Consulte esta [colección](https://huggingface.co/collections/Photoroom/prx-models-68e66254c202ebfab99ad38e) para obtener más información.

## Carga del pipeline

Load the pipeline with [from_pretrained()](/docs/diffusers/main/es/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")
```

### Carga manual de componentes

Cargue los componentes de forma individual para personalizar el pipeline, por ejemplo, para usar modelos cuantizados.

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

## Optimización de memoria

Para entornos con limitaciones de memoria:

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

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

**Parámetros:**

transformer (`PRXTransformer2DModel`) : El modelo transformer de PRX para desruidar los latentes de imagen codificados.

scheduler ([FlowMatchEulerDiscreteScheduler](/docs/diffusers/main/es/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler)) : Un programador para usar en combinación con `transformer` para desruidar los latentes de imagen codificados.

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

tokenizer ([`T5TokenizerFast` o `GemmaTokenizerFast`]) : Tokenizador para el codificador de texto.

vae ([AutoencoderKL](/docs/diffusers/main/es/api/models/autoencoderkl#diffusers.AutoencoderKL) o [AutoencoderDC](/docs/diffusers/main/es/api/models/autoencoder_dc#diffusers.AutoencoderDC)) : Modelo de Autoencoder Variacional (VAE) para codificar y decodificar imágenes hacia y desde representaciones latentes. Admite tanto AutoencoderKL (compresión 8x) como AutoencoderDC (compresión 32x).

Pipeline para la generación de imágenes a partir de texto usando PRX Transformer.

Este modelo hereda de [DiffusionPipeline](/docs/diffusers/main/es/api/pipelines/overview#diffusers.DiffusionPipeline). Consulta la documentación de la clase superior para los métodos genéricos que la biblioteca implementa para todos los pipelines (como la descarga o el guardado, la ejecución en un 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 = 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 = ['latents'], tokenizer_max_length: int | None = None, skip_text_cleaning: bool = False)
```

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

**Parámetros:**

prompt (*str* o *list[str]*, *opcional*) : El prompt o los prompts para guiar la generación de la imagen. Si no se define, se debe pasar *prompt_embeds* en su lugar.

negative_prompt (*str*, *opcional*, valor predeterminado *""*) : El prompt o prompts que no deben guiar la generación de la imagen. Se ignora cuando no se usa guía (es decir, se ignora si *guidance_scale* es menor que *1*).

height (*int*, *optional*, por defecto self.transformer.config.sample_size * self.vae_scale_factor) : La altura en píxeles de la imagen generada.

width (*int*, *optional*, por defecto self.transformer.config.sample_size * self.vae_scale_factor) : El ancho en píxeles de la imagen generada.

num_inference_steps (*int*, *optional*, por defecto 28) : El número de pasos de desruido. Más pasos de desruido suelen conducir a una imagen de mayor calidad a costa de una inferencia más lenta.

timesteps (*list[int]*, *optional*) : Pasos de tiempo personalizados para usar en el proceso de desruido con programadores que admiten un argumento *timesteps* en su método *set_timesteps*. Si no se define, se usará el comportamiento predeterminado cuando se pase *num_inference_steps*. Debe estar en orden descendente.

guidance_scale (*float*, *optional*, valor predeterminado 4.0) : Escala de orientación definida en [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598). *guidance_scale* se define como *w* de la ecuación 2. de [Imagen Paper](https://huggingface.co/papers/2205.11487). La escala de orientación se habilita estableciendo *guidance_scale > 1*. Una escala de orientación más alta fomenta la generación de imágenes estrechamente vinculadas al *prompt* de texto, generalmente a costa de una menor calidad de imagen.

num_images_per_prompt (*int*, *optional*, por defecto 1) : El número de imágenes a generar por prompt.

generator (*torch.Generator* o *list[torch.Generator]*, *opcional*) : Uno o una lista de [generadores de torch](https://pytorch.org/docs/stable/generated/torch.Generator.html) para hacer la generación determinista.

latents (*torch.Tensor*, *opcional*) : Latentes ruidosos pregenerados, muestreados de una distribución gaussiana, que se utilizarán como entradas para la generación de imágenes. Se pueden usar para ajustar la misma generación con diferentes prompts. Si no se proporcionan, se generará un tensor de latentes mediante muestreo usando el *generator* aleatorio suministrado.

prompt_embeds (*torch.FloatTensor*, *opcional*) : Incrustaciones de texto pregeneradas. Se pueden utilizar para ajustar fácilmente las entradas de texto, *por ejemplo*, el ponderado de prompts. Si no se proporcionan, las incrustaciones de texto se generarán a partir del argumento de entrada *prompt*.

negative_prompt_embeds (*torch.FloatTensor*, *opcional*) : Embeddings de texto negativos pregenerados. Se pueden utilizar para ajustar fácilmente las entradas de texto, *p. ej.* el ponderado de prompts. Si no se proporcionan y *guidance_scale > 1*, se generarán embeddings negativos a partir de una cadena vacía.

prompt_attention_mask (*torch.BoolTensor*, *opcional*) : Máscara de atención pregenerada para *prompt_embeds*. Si no se proporciona, la máscara de atención se generará a partir del argumento de entrada *prompt*.

negative_prompt_attention_mask (*torch.BoolTensor*, *opcional*) : Máscara de atención pregenerada para *negative_prompt_embeds*. Si no se proporciona y *guidance_scale > 1*, la máscara de atención se generará a partir de una cadena vacía.

output_type (*str*, *opcional*, por defecto *"pil"*) : El formato de salida de la imagen generada. Elija entre [PIL](https://pillow.readthedocs.io/en/stable/): *PIL.Image.Image* o *np.array*.

tokenizer_max_length (*int*, *opcional*) : Anula el número máximo de tokens utilizados al tokenizar el prompt. Si no se establece, el valor predeterminado es el `model_max_length` propio del tokenizador.

skip_text_cleaning (*bool*, *opcional*, valor predeterminado *False*) : Si es *True*, utiliza solo una limpieza ligera del prompt (corregir codificación + desescapar HTML) en lugar del pipeline completo de limpieza de DeepFloyd.

return_dict (*bool*, *opcional*, valor predeterminado *True*) : Indica si se debe devolver una [*~pipelines.prx.PRXPipelineOutput*] en lugar de una tupla simple.

use_resolution_binning (*bool*, *opcional*, valor predeterminado *True*) : Si se establece en *True*, la altura y el ancho solicitados se mapean primero a las resoluciones más cercanas utilizando bins de relación de aspecto predefinidos. Después de que los latentes generados se decodifican en imágenes, se redimensionan de nuevo a la resolución solicitada. Útil para generar imágenes no cuadradas en resoluciones óptimas.

callback_on_step_end (*Callable*, *opcional*) : Una función que se llama al final de cada paso de desruido durante la inferencia. La función se llama con los siguientes argumentos: *callback_on_step_end(self, step, timestep, callback_kwargs)*. *callback_kwargs* incluirá una lista de todos los tensores especificados por *callback_on_step_end_tensor_inputs*.

callback_on_step_end_tensor_inputs (*list*, *optional*) : La lista de entradas de tensores para la función *callback_on_step_end*. Los tensores especificados en la lista se pasarán como argumento *callback_kwargs*. Solo podrá incluir tensores que estén listados en el atributo *._callback_tensor_inputs*.

**Devuelve:** [*~pipelines.prx.PRXPipelineOutput*] o *tuple`

[*~pipelines.prx.PRXPipelineOutput*] si *return_dict* es True, de lo contrario una *tupla. Al devolver una tupla, el primer elemento es una lista con las imágenes generadas.

Función invocada al llamar a la pipeline para la generación.

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

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

Comprueba que todas las entradas estén en el formato correcto.

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

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

Codifica el prompt de texto usando el codificador y el tokenizador de texto estándar, o usa embeddings precalculados.

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

```python
get_default_resolution()
```

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

**Devuelve:** `int`

Tamaño de muestra predeterminado (alto/ancho) para usar en la generación.

Determine la resolución predeterminada en función del VAE y la configuración cargados.

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

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

Prepara los latentes iniciales para el proceso de difusión.

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

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

**Parámetros:**

images (`list[PIL.Image.Image]` o `np.ndarray`) : lista de imágenes PIL desruidas de longitud `batch_size` o matriz numpy de forma `(batch_size, height, width, num_channels)`. Las imágenes PIL o la matriz numpy presentan las imágenes desruidas del pipeline de difusión.

Clase de salida para las tuberías de PRX.