Diffusers documentation
PRX
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 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 | 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 | 512 | Sim | Não | Ajuste fino realizado no conjunto de dados Alchemist dataset com Flux VAE | Consegue lidar com prompts menos detalhados | 28 etapas, cfg=5.0 | torch.bfloat16 |
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 | 512 | Sim | Não | Ajuste fino realizado no conjunto de dados Alchemist dataset 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 | 512 | Sim | Sim | Modelo destilado de 8 etapas de 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 | 512 | Não | Não | Modelo base pré-treinado em 512 com Deep Compression Autoencoder (DC-AE) | Funciona melhor com prompts detalhados em linguagem natural | 28 etapas, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae-sft | 512 | Sim | Não | Ajustado no conjunto de dados Alchemist dataset com Deep Compression Autoencoder (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 | 512 | Sim | Sim | Modelo destilado de 8 etapas de Photoroom/prx-512-t2i-dc-ae-sft-distilled | Consegue lidar com prompts menos detalhados em linguagem natural | 8 etapas, cfg=1.0 | torch.bfloat16 |
Consulte esta coleção para mais informações.
Carregando o 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")Carregamento Manual de Componentes
Carregue os componentes individualmente para personalizar o pipeline, por exemplo, para usar modelos quantizados.
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:
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 )
Parâmetros
- transformer (
PRXTransformer2DModel) — O modelo transformer PRX para desruidar os latentes da imagem codificada. - scheduler (FlowMatchEulerDiscreteScheduler) —
Um scheduler a ser usado em combinação com
transformerpara desruidar os latentes da imagem codificada. - text_encoder (
T5GemmaEncoder) — Modelo de codificador de texto para codificar prompts. - tokenizer ([
T5TokenizerFastouGemmaTokenizerFast]) — Tokenizer para o codificador de texto. - vae (AutoencoderKL ou AutoencoderDC) — Modelo 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. 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__
< 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`
Parâmetros
- prompt (str ou list[str], opcional) — O prompt ou os prompts para orientar a geração da imagem. Se não for 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, opcional, padrão self.transformer.config.sample_size * self.vae_scale_factor) — A altura em pixels da imagem gerada.
- width (int, opcional, padrão é self.transformer.config.sample_size * self.vae_scale_factor) — A largura em pixels da imagem gerada.
- num_inference_steps (int, opcional, 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], opcional) — Timesteps 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, opcional, padrão 4.0) — Escala de orientação definida em Classifier-Free Diffusion Guidance. guidance_scale é definido como w da equação 2 do Imagen Paper. 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, opcional, 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 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) — Incorporações de texto pré-geradas. Podem ser usadas para ajustar facilmente as entradas de texto, por exemplo, ponderação de prompt. Se não forem fornecidas, as incorporações de texto serão geradas a partir do argumento de entrada prompt.
- negative_prompt_embeds (torch.FloatTensor, opcional) — Incorporações de texto negativas pré-geradas. 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, as incorporações negativas serão geradas 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: PIL.Image.Image ou np.array.
- tokenizer_max_length (int, opcional) —
Substitui o número máximo de tokens usados ao tokenizar o prompt. O padrão é o próprio
model_max_lengthdo tokenizer, quando não definido. - skip_text_cleaning (bool, opcional, padrão False) — Se True, usa apenas 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:
>>> 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 )
Verifique se todas as entradas estão no formato correto.
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 )
Codifique o prompt de texto usando um codificador e um tokenizador de texto padrão, ou use embeddings pré-computados.
get_default_resolution
< source >( ) → int
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
< source >( batch_size: intnum_channels_latents: intheight: intwidth: intdtype: dtypedevice: devicegenerator: typing.Optional[torch.Generator] = Nonelatents: typing.Optional[torch.Tensor] = None )
Prepare os latentes iniciais para o processo de difusão.
PRXPipelineOutput
classe diffusers.pipelines.prx.PRXPipelineOutput
< source >( images: list[PIL.Image.Image] | numpy.ndarray )
Classe de saída para pipelines PRX.