#
# Sous licence Apache License, Version 2.0 (la « Licence ») ;
# vous ne pouvez utiliser ce fichier que conformément à la Licence.
# Vous pouvez obtenir un exemplaire de la Licence à
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Sauf exigence légale ou accord écrit
# distribué sous la Licence est distribué sur une base « EN L'ÉTAT »,
# SANS AUCUNE GARANTIE NI CONDITION, explicite ou implicite.
# Voir la Licence pour le langage spécifique régissant les autorisations et
# limitations under the License. -->

# PRX

PRX génère des images de haute qualité à partir de texte en utilisant une architecture MMDIT simplifiée où les jetons de texte ne sont pas mis à jour à travers les blocs de transformeur. Il emploie le flow matching avec un planificateur discret pour un échantillonnage efficace et utilise le modèle T5Gemma-2B-2B-UL2 de Google pour l'encodage de texte multilingue. Le transformeur d'environ 1,3 milliard de paramètres offre une inférence rapide sans sacrifier la qualité. Vous pouvez choisir entre Flux VAE (compression 8x, 16 canaux latents) pour un équilibre entre qualité et vitesse, ou DC-AE (compression 32x, 32 canaux latents) pour une compression latente et un traitement plus rapide.

## Modèles disponibles

PRX propose plusieurs variantes avec des configurations VAE différentes, chacune optimisée pour des résolutions spécifiques. Les modèles de base excellent avec des prompts détaillés, capturant des compositions complexes et des détails subtils. Les modèles affinés entraînés sur le [jeu de données Alchemist](https://huggingface.co/datasets/yandex/alchemist) améliorent la qualité esthétique, en particulier avec des prompts plus simples.

| Modèle | Résolution | Ajusté | Distillé | Description | Prompts suggérés | Paramètres suggérés | dtype recommandé |
|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|
| [`Photoroom/prx-256-t2i`](https://huggingface.co/Photoroom/prx-256-t2i)| 256 | Non | Non | Modèle de base pré-entraîné à 256 avec le VAE de Flux|Fonctionne le mieux avec des prompts détaillés en langage naturel|28 étapes, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-256-t2i-sft`](https://huggingface.co/Photoroom/prx-256-t2i-sft)| 512 | Oui | Non | Ajusté finement sur le jeu de données [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) avec Flux VAE | Peut traiter des prompts moins détaillés|28 étapes, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i`](https://huggingface.co/Photoroom/prx-512-t2i)| 512 | Non | Non | Modèle de base pré-entraîné à 512 avec le VAE de Flux |Fonctionne le mieux avec des prompts détaillés en langage naturel|28 étapes, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft)| 512 | Oui | Non | Ajusté finement sur le jeu de données [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) avec Flux VAE | Capable de traiter des prompts moins détaillés en langage naturel|28 étapes, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-sft-distilled)| 512 | Oui | Oui | Modèle distillé en 8 étapes issu de [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft) | Capable de traiter des prompts moins détaillés en langage naturel|8 étapes, cfg=1.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae)| 512 | Non | Non | Modèle de base pré-entraîné à 512 avec [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)|Fonctionne le mieux avec des prompts détaillés en langage naturel|28 étapes, cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft)| 512 | Oui | Non | Ajusté finement sur le jeu de données [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) avec [Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae) | Capable de traiter des prompts moins détaillés en langage naturel|28 étapes, 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 | Oui | Oui | Modèle distillé en 8 étapes issu de [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled) | Peut traiter des prompts moins détaillés en langage naturel|8 étapes, cfg=1.0| `torch.bfloat16` |s

Consultez cette [collection](https://huggingface.co/collections/Photoroom/prx-models-68e66254c202ebfab99ad38e) pour plus d'informations.

## Chargement du pipeline

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

### Chargement manuel des composants

Chargez les composants individuellement pour personnaliser le pipeline, par exemple pour utiliser des modèles quantifiés.

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

## Optimisation de la mémoire

Pour les environnements à mémoire limitée :

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

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

**Paramètres :**

transformer (`PRXTransformer2DModel`) : Le modèle transformer PRX pour le débruitage des latents d'image encodés.

scheduler ([FlowMatchEulerDiscreteScheduler](/docs/diffusers/main/fr/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler)) : un planificateur à utiliser en combinaison avec `transformer` pour débruyer les latents d'image encodés.

text_encoder (`T5GemmaEncoder`) : Modèle d'encodeur de texte pour l'encodage des invites.

tokenizer ([`T5TokenizerFast` ou `GemmaTokenizerFast`]) : Tokenizer pour l'encodeur de texte.

vae ([AutoencoderKL](/docs/diffusers/main/fr/api/models/autoencoderkl#diffusers.AutoencoderKL) ou [AutoencoderDC](/docs/diffusers/main/fr/api/models/autoencoder_dc#diffusers.AutoencoderDC)) : modèle de Variational Auto-Encoder (VAE) pour encoder et décoder des images vers et depuis des représentations latentes. Prend en charge AutoencoderKL (compression 8x) et AutoencoderDC (compression 32x).

Pipeline pour la génération d'images à partir de texte utilisant PRX Transformer.

Ce modèle hérite de [DiffusionPipeline](/docs/diffusers/main/fr/api/pipelines/overview#diffusers.DiffusionPipeline). Consultez la documentation de la classe parente pour les méthodes génériques que la bibliothèque implémente pour tous les pipelines (téléchargement ou sauvegarde, exécution sur un appareil particulier, etc.).

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

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

**Paramètres :**

prompt (*str* ou *list[str]*, *facultatif*) : Le prompt ou les prompts pour guider la génération d’image. S’il n’est pas défini, il faut passer *prompt_embeds* à la place.

negative_prompt (*str*, *facultatif*, par défaut *""*) : Le prompt ou les prompts à ne pas utiliser pour guider la génération d'image. Ignoré lorsque la guidance n'est pas utilisée (c'est-à-dire ignoré si *guidance_scale* est inférieur à *1*).

height (*int*, *optionnel*, par défaut self.transformer.config.sample_size * self.vae_scale_factor) : La hauteur en pixels de l'image générée.

width (*int*, *optionnel*, par défaut self.transformer.config.sample_size * self.vae_scale_factor) : La largeur en pixels de l'image générée.

num_inference_steps (*int*, *optionnel*, par défaut 28) : Le nombre d'étapes de débruitage. Un plus grand nombre d'étapes de débruitage conduit généralement à une image de meilleure qualité, au prix d'une inférence plus lente.

timesteps (*list[int]*, *optionnel*) : Étapes de temps personnalisées à utiliser pour le processus de débruitage avec les planificateurs qui prennent en charge un argument *timesteps* dans leur méthode *set_timesteps*. S'il n'est pas défini, le comportement par défaut lorsque *num_inference_steps* est passé sera utilisé. Doit être dans l'ordre décroissant.

guidance_scale (*float*, *optional*, par défaut 4.0) : Échelle de guidage telle que définie dans [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598). *guidance_scale* est défini comme *w* de l'équation 2. de [Imagen Paper](https://huggingface.co/papers/2205.11487). L'échelle de guidage est activée en définissant *guidance_scale > 1*. Une échelle de guidage plus élevée encourage la génération d'images étroitement liées au *prompt* texte, généralement au détriment de la qualité de l'image.

num_images_per_prompt (*int*, *optionnel*, valeur par défaut 1) : Le nombre d'images à générer par invite.

generator (*torch.Generator* ou *list[torch.Generator]*, *facultatif*) : Un ou une liste de [générateurs torch](https://pytorch.org/docs/stable/generated/torch.Generator.html) pour rendre la génération déterministe.

latents (*torch.Tensor*, *optionnel*) : Latents bruités pré-générés, échantillonnés à partir d'une distribution gaussienne, à utiliser comme entrées pour la génération d'images. Peut être utilisé pour ajuster la même génération avec des prompts différents. S'il n'est pas fourni, un tenseur de latents sera généré par échantillonnage à l'aide du *generator* aléatoire fourni.

prompt_embeds (*torch.FloatTensor*, *facultatif*) : Incrustations de texte pré-générées. Peut être utilisé pour ajuster facilement les entrées de texte, *par exemple* le pondération des prompts. Si non fourni, les incrustations de texte seront générées à partir de l'argument d'entrée *prompt*.

negative_prompt_embeds (*torch.FloatTensor*, *facultatif*) : Embeddings de texte négatifs pré-générés. Peut être utilisé pour ajuster facilement les entrées de texte, *par exemple* le pondération des prompts. S'il n'est pas fourni et que *guidance_scale > 1*, les embeddings négatifs seront générés à partir d'une chaîne vide.

prompt_attention_mask (*torch.BoolTensor*, *facultatif*) : Masque d'attention pré-généré pour *prompt_embeds*. S'il n'est pas fourni, le masque d'attention sera généré à partir de l'argument d'entrée *prompt*.

negative_prompt_attention_mask (*torch.BoolTensor*, *facultatif*) : Masque d'attention pré-généré pour *negative_prompt_embeds*. S'il n'est pas fourni et que *guidance_scale > 1*, le masque d'attention sera généré à partir d'une chaîne vide.

output_type (*str*, *optionnel*, par défaut *"pil"*) : Le format de sortie de l'image générée. Choisissez entre [PIL](https://pillow.readthedocs.io/en/stable/) : *PIL.Image.Image* ou *np.array*.

tokenizer_max_length (*int*, *facultatif*) : Remplace le nombre maximal de jetons utilisés lors de la tokenisation du prompt. Par défaut, utilise le `model_max_length` du tokenizer lui-même si non défini.

skip_text_cleaning (*bool*, *facultatif*, par défaut *False*) : Si *True*, n'utilise qu'un nettoyage léger du prompt (correction de l'encodage + déséchappement HTML) au lieu du pipeline de nettoyage complet de DeepFloyd.

return_dict (*bool*, *facultatif*, par défaut *True*) : Indique s'il faut retourner un [*~pipelines.prx.PRXPipelineOutput*] plutôt qu'un simple tuple.

use_resolution_binning (*bool*, *facultatif*, valeur par défaut *True*) : Si défini sur *True*, la hauteur et la largeur demandées sont d’abord mappées sur les résolutions les plus proches à l’aide de bins de ratio d’aspect prédéfinis. Après le décodage des latents produits en images, celles-ci sont redimensionnées pour revenir à la résolution demandée. Utile pour générer des images non carrées à des résolutions optimales.

callback_on_step_end (*Callable*, *facultatif*) : Une fonction appelée à la fin de chaque étape de débruitage pendant l’inférence. La fonction est appelée avec les arguments suivants : *callback_on_step_end(self, step, timestep, callback_kwargs)*. *callback_kwargs* inclura une liste de tous les tenseurs, tels que spécifiés par *callback_on_step_end_tensor_inputs*.

callback_on_step_end_tensor_inputs (*list*, *optionnel*) : La liste des entrées tensorielles pour la fonction *callback_on_step_end*. Les tenseurs spécifiés dans la liste seront passés en tant qu'argument *callback_kwargs*. Vous ne pourrez inclure que les tenseurs répertoriés dans l'attribut *._callback_tensor_inputs*.

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

[*~pipelines.prx.PRXPipelineOutput*] si *return_dict* est True, sinon un *tuple. En cas de retour d'un tuple, le premier élément est une liste contenant les images générées.

Fonction appelée lors de l'appel du pipeline pour la génération.

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

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

Vérifiez que toutes les entrées sont au format correct.

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

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

Encodez le prompt de texte à l'aide d'un encodeur et d'un tokeniseur de texte standard, ou utilisez des embeddings précalculés.

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

```python
get_default_resolution()
```

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

**Retourne :** `int`

Taille d'échantillon par défaut (hauteur/largeur) à utiliser pour la génération.

Déterminer la résolution par défaut en fonction du VAE chargé et de la configuration.

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

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

Préparez les latents initiaux pour le processus de diffusion.

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

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

**Paramètres :**

images (`list[PIL.Image.Image]` ou `np.ndarray`) : liste d'images PIL débruitées de longueur `batch_size` ou tableau numpy de forme `(batch_size, height, width, num_channels)`. Les images PIL ou le tableau numpy présentent les images débruitées du pipeline de diffusion.

Classe de sortie pour les pipelines PRX.

