Diffusers documentation
PRX
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 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 | 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 | 512 | Oui | Non | Ajusté finement sur le jeu de données Alchemist dataset avec Flux VAE | Peut traiter des prompts moins détaillés | 28 étapes, cfg=5.0 | torch.bfloat16 |
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 | 512 | Oui | Non | Ajusté finement sur le jeu de données Alchemist dataset 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 | 512 | Oui | Oui | Modèle distillé en 8 étapes issu de 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 | 512 | Non | Non | Modèle de base pré-entraîné à 512 avec Deep Compression Autoencoder (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 | 512 | Oui | Non | Ajusté finement sur le jeu de données Alchemist dataset avec Deep Compression Autoencoder (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 | 512 | Oui | Oui | Modèle distillé en 8 étapes issu de Photoroom/prx-512-t2i-dc-ae-sft-distilled | Capable de traiter des prompts moins détaillés en langage naturel | 8 étapes, cfg=1.0 | torch.bfloat16 |
Consultez cette collection pour plus d'informations.
Chargement du pipeline
Load the pipeline with from_pretrained().
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.
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 :
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
classe diffusers.PRXPipeline
< source >( transformer: PRXTransformer2DModelscheduler: FlowMatchEulerDiscreteSchedulertext_encoder: T5GemmaEncodertokenizer: transformers.models.t5.tokenization_t5.T5Tokenizer | transformers.models.gemma.tokenization_gemma.GemmaTokenizer | transformers.models.auto.tokenization_auto.AutoTokenizervae: diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL | diffusers.models.autoencoders.autoencoder_dc.AutoencoderDC | None = Nonedefault_sample_size: int | None = 512 )
Paramètres
- transformer (
PRXTransformer2DModel) — Le modèle transformer PRX pour débruyer les latents d'image encodés. - scheduler (FlowMatchEulerDiscreteScheduler) —
Un planificateur à utiliser en combinaison avec
transformerpour débruyer les latents d'image encodés. - text_encoder (
T5GemmaEncoder) — Modèle d'encodeur de texte pour encoder les prompts. - tokenizer ([
T5TokenizerFastouGemmaTokenizerFast]) — Jetoniseur pour l'encodeur de texte. - vae (AutoencoderKL ou AutoencoderDC) — Modèle 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. Consultez la documentation de la classe parente pour les méthodes génériques que la bibliothèque implémente pour tous les pipelines (tels que le téléchargement ou l'enregistrement, l'exécution sur un appareil particulier, etc.)
__call__
< source >( prompt: str | list[str] = Nonenegative_prompt: str = ''height: int | None = Nonewidth: int | None = Nonenum_inference_steps: int = 28timesteps: list = Noneguidance_scale: float = 4.0num_images_per_prompt: int | None = 1generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneprompt_embeds: typing.Optional[torch.FloatTensor] = Nonenegative_prompt_embeds: typing.Optional[torch.FloatTensor] = Noneprompt_attention_mask: typing.Optional[torch.BoolTensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.BoolTensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Trueuse_resolution_binning: bool = Truecallback_on_step_end: typing.Optional[typing.Callable[[int, int], NoneType]] = Nonecallback_on_step_end_tensor_inputs: list = ['latents']tokenizer_max_length: int | None = Noneskip_text_cleaning: bool = False ) → [~pipelines.prx.PRXPipelineOutput] or *tuple`
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 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, facultatif, 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, facultatif, 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, facultatif, valeur 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], facultatif) — Timesteps personnalisés à utiliser pour le processus de débruitage avec les planificateurs qui prennent en charge un argument timesteps dans leur méthode set_timesteps. S'ils ne sont pas définis, le comportement par défaut lorsque num_inference_steps est passé sera utilisé. Doivent être dans l'ordre décroissant.
- guidance_scale (float, facultatif, par défaut 4.0) — Échelle de guidage telle que définie dans Classifier-Free Diffusion Guidance. guidance_scale est défini comme w de l'équation 2 de Imagen Paper. 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, facultatif, 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 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) — Embeddings de texte pré-générés. Peut être utilisé pour ajuster facilement les entrées de texte, par ex. le pondération des prompts. S'il n'est pas fourni, les embeddings de texte seront générés à 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 ex. 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, facultatif, par défaut “pil”) — Le format de sortie de l'image générée. Choisissez entre PIL : 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, la valeur est celle de
model_max_lengthdu tokenizer si elle n’est pas définie. - skip_text_cleaning (bool, facultatif, valeur 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, comme spécifié 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 des 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 :
>>> 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
< source >( prompt: str | list[str]height: intwidth: intguidance_scale: floatcallback_on_step_end_tensor_inputs: list[str] | None = Noneprompt_embeds: typing.Optional[torch.FloatTensor] = Nonenegative_prompt_embeds: typing.Optional[torch.FloatTensor] = None )
Vérifiez que toutes les entrées sont au format correct.
encode_prompt
< source >( prompt: str | list[str]device: typing.Optional[torch.device] = Nonedo_classifier_free_guidance: bool = Truenegative_prompt: str = ''num_images_per_prompt: int = 1prompt_embeds: typing.Optional[torch.FloatTensor] = Nonenegative_prompt_embeds: typing.Optional[torch.FloatTensor] = Noneprompt_attention_mask: typing.Optional[torch.BoolTensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.BoolTensor] = Nonetokenizer_max_length: int | None = Noneskip_text_cleaning: bool = False )
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
< source >( ) → int
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
< source >( batch_size: intnum_channels_latents: intheight: intwidth: intdtype: dtypedevice: devicegenerator: typing.Optional[torch.Generator] = Nonelatents: typing.Optional[torch.Tensor] = None )
Préparez les latents initiaux pour le processus de diffusion.
PRXPipelineOutput
classe diffusers.pipelines.prx.PRXPipelineOutput
< source >( images: list[PIL.Image.Image] | numpy.ndarray )
Classe de sortie pour les pipelines PRX.