Diffusers documentation
PRX
PRX
PRX tạo ra các hình ảnh chất lượng cao từ văn bản bằng cách sử dụng kiến trúc MMDIT đơn giản hóa, trong đó các token văn bản không được cập nhật thông qua các khối transformer. Nó áp dụng flow matching với lịch trình rời rạc để lấy mẫu hiệu quả và sử dụng mô hình T5Gemma-2B-2B-UL2 của Google để mã hóa văn bản đa ngôn ngữ. Transformer với khoảng 1,3 tỷ tham số này cung cấp suy luận nhanh mà không làm giảm chất lượng. Bạn có thể chọn giữa Flux VAE (nén 8 lần, 16 kênh tiềm ẩn) để cân bằng giữa chất lượng và tốc độ, hoặc DC-AE (nén 32 lần, 32 kênh tiềm ẩn) để nén tiềm ẩn và xử lý nhanh hơn.
Các model khả dụng
PRX cung cấp nhiều biến thể với các cấu hình VAE khác nhau, mỗi biến thể được tối ưu hóa cho các độ phân giải cụ thể. Các mô hình cơ sở thể hiện tốt với các prompt chi tiết, nắm bắt được các bố cục phức tạp và những chi tiết tinh tế. Các mô hình được tinh chỉnh huấn luyện trên dataset Alchemist cải thiện chất lượng thẩm mỹ, đặc biệt là với các prompt đơn giản hơn.
| Mô hình | Giải pháp | Được tinh chỉnh | Được chưng cất | Mô tả | Gợi ý các câu lệnh | Các tham số được đề xuất | dtype được khuyến nghị |
|---|---|---|---|---|---|---|---|
Photoroom/prx-256-t2i | 256 | Không | Không | Mô hình cơ sở được huấn luyện trước ở 256 với Flux VAE | Hoạt động tốt nhất với các prompt chi tiết bằng ngôn ngữ tự nhiên | 28 bước, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-256-t2i-sft | 512 | Có | Không | Được tinh chỉnh trên bộ dữ liệu Alchemist dataset với Flux VAE | Có thể xử lý các prompt ít chi tiết hơn | 28 bước, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i | 512 | Không | Không | Mô hình cơ sở được huấn luyện trước ở 512 với Flux VAE | Hoạt động tốt nhất với các prompt chi tiết bằng ngôn ngữ tự nhiên | 28 bước, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-sft | 512 | Có | Không | Được tinh chỉnh trên bộ dữ liệu Alchemist dataset với Flux VAE | Có thể xử lý các prompt ít chi tiết hơn bằng ngôn ngữ tự nhiên | 28 bước, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-sft-distilled | 512 | Có | Có | Mô hình chưng cất 8 bước từ Photoroom/prx-512-t2i-sft | Có thể xử lý các prompt ít chi tiết hơn bằng ngôn ngữ tự nhiên | 8 bước, cfg=1.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae | 512 | Không | Không | Mô hình cơ sở được huấn luyện trước ở 512 với Deep Compression Autoencoder (DC-AE) | Hoạt động tốt nhất với các prompt chi tiết bằng ngôn ngữ tự nhiên | 28 bước, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae-sft | 512 | Có | Không | Được tinh chỉnh trên bộ dữ liệu Alchemist dataset với Deep Compression Autoencoder (DC-AE) | Có thể xử lý các prompt ít chi tiết hơn bằng ngôn ngữ tự nhiên | 28 bước, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae-sft-distilled | 512 | Có | Có | Mô hình chưng cất 8 bước từ Photoroom/prx-512-t2i-dc-ae-sft-distilled | Có thể xử lý các prompt ít chi tiết hơn bằng ngôn ngữ tự nhiên | 8 bước, cfg=1.0 | torch.bfloat16 |
Tham khảo bộ sưu tập này để biết thêm thông tin.
Tải pipeline
Tải pipeline bằng 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")Tải thành phần thủ công
Tải các thành phần riêng lẻ để tùy chỉnh pipeline, chẳng hạn như để sử dụng các mô hình đã lượng tử hóa.
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"Tối ưu hóa bộ nhớ
Đối với các môi trường bị giới hạn bộ nhớ:
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
lớp 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 )
Tham số
- transformer (
PRXTransformer2DModel) — Mô hình transformer PRX để khử nhiễu các latents của ảnh đã mã hóa. - scheduler (FlowMatchEulerDiscreteScheduler) —
Một scheduler được sử dụng kết hợp với
transformerđể khử nhiễu các latents ảnh đã được mã hóa. - text_encoder (
T5GemmaEncoder) — Mô hình mã hóa văn bản để mã hóa các prompt. - tokenizer ([
T5TokenizerFasthoặcGemmaTokenizerFast]) — Tokenizer cho text encoder. - vae (AutoencoderKL hoặc AutoencoderDC) — Mô hình Variational Auto-Encoder (VAE) để mã hóa và giải mã hình ảnh từ và sang các biểu diễn ẩn. Hỗ trợ cả AutoencoderKL (nén 8 lần) và AutoencoderDC (nén 32 lần).
Pipeline để tạo ảnh từ văn bản sử dụng PRX Transformer.
Mô hình này kế thừa từ DiffusionPipeline. Xem tài liệu của lớp cha để biết các phương thức chung mà thư viện triển khai cho tất cả các pipeline (chẳng hạn như tải xuống hoặc lưu, chạy trên một thiết bị cụ thể, v.v.)
__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`
Tham số
- prompt (str hoặc list[str], tuỳ chọn) — Prompt hoặc các prompt để định hướng việc tạo ảnh. Nếu không được định nghĩa, bạn phải truyền prompt_embeds thay thế.
- negative_prompt (str, tuỳ chọn, mặc định là "") — Prompt hoặc các prompt không dùng để định hướng việc tạo ảnh. Bỏ qua khi không sử dụng guidance (tức là bị bỏ qua nếu guidance_scale nhỏ hơn 1).
- height (int, optional, mặc định là self.transformer.config.sample_size * self.vae_scale_factor) — Chiều cao của hình ảnh được tạo ra, tính bằng pixel.
- width (int, tuỳ chọn, mặc định là self.transformer.config.sample_size * self.vae_scale_factor) — Chiều rộng của hình ảnh được tạo ra, tính bằng pixel.
- num_inference_steps (int, tuỳ chọn, mặc định là 28) — Số bước khử nhiễu. Số bước khử nhiễu nhiều hơn thường dẫn đến hình ảnh có chất lượng cao hơn với chi phí là suy luận chậm hơn.
- timesteps (list[int], optional) — Các bước thời gian (timesteps) tùy chỉnh để sử dụng cho quá trình khử nhiễu với các trình lập lịch (scheduler) hỗ trợ tham số timesteps trong phương thức set_timesteps của chúng. Nếu không được định nghĩa, hành vi mặc định khi num_inference_steps được truyền vào sẽ được sử dụng. Phải được sắp xếp theo thứ tự giảm dần.
- guidance_scale (float, optional, mặc định là 4.0) — Thang độ dẫn hướng như được định nghĩa trong Classifier-Free Diffusion Guidance. guidance_scale được định nghĩa là w của phương trình 2. trong Imagen Paper. Thang độ dẫn hướng được kích hoạt bằng cách đặt guidance_scale > 1. Thang độ dẫn hướng cao hơn khuyến khích tạo ra các hình ảnh gắn kết chặt chẽ với prompt văn bản, thường phải đánh đổi bằng chất lượng hình ảnh thấp hơn.
- num_images_per_prompt (int, tuỳ chọn, mặc định là 1) — Số lượng hình ảnh cần tạo cho mỗi prompt.
- generator (torch.Generator hoặc list[torch.Generator], tùy chọn) — Một hoặc một danh sách torch generator(s) để đảm bảo quá trình sinh là tất định.
- latents (torch.Tensor, optional) — Các latents nhiễu đã được tạo trước, lấy mẫu từ phân phối Gaussian, dùng làm đầu vào cho việc tạo ảnh. Có thể dùng để tinh chỉnh cùng một quá trình tạo với các prompt khác nhau. Nếu không được cung cấp, một tensor latents sẽ được tạo bằng cách lấy mẫu sử dụng generator ngẫu nhiên được cung cấp.
- prompt_embeds (torch.FloatTensor, optional) — Các embedding văn bản đã được tạo trước. Có thể được sử dụng để dễ dàng điều chỉnh các đầu vào văn bản, vd. trọng số prompt. Nếu không được cung cấp, các embedding văn bản sẽ được tạo từ đối số đầu vào prompt.
- negative_prompt_embeds (torch.FloatTensor, optional) — Các embedding văn bản âm đã được tạo trước. Có thể được sử dụng để dễ dàng tinh chỉnh các đầu vào văn bản, vd. trọng số prompt. Nếu không được cung cấp và guidance_scale > 1, các embedding âm sẽ được tạo từ một chuỗi rỗng.
- prompt_attention_mask (torch.BoolTensor, optional) — Mặt nạ chú ý (attention mask) được tạo trước cho prompt_embeds. Nếu không được cung cấp, mặt nạ chú ý sẽ được tạo từ đối số đầu vào prompt.
- negative_prompt_attention_mask (torch.BoolTensor, optional) — Attention mask được tạo sẵn cho negative_prompt_embeds. Nếu không được cung cấp và guidance_scale > 1, attention mask sẽ được tạo từ một chuỗi rỗng.
- output_type (str, optional, mặc định là “pil”) — Định dạng đầu ra của hình ảnh được tạo. Chọn giữa PIL: PIL.Image.Image hoặc np.array.
- tokenizer_max_length (int, optional) —
Ghi đè số lượng token tối đa được sử dụng khi thực hiện token hóa prompt. Mặc định là
model_max_lengthcủa tokenizer khi không được đặt. - skip_text_cleaning (bool, optional, mặc định là False) — Nếu True, chỉ sử dụng làm sạch prompt nhẹ (sửa mã hóa + bỏ escape HTML) thay vì toàn bộ pipeline làm sạch DeepFloyd.
- return_dict (bool, tuỳ chọn, mặc định là True) — Có nên trả về [~pipelines.prx.PRXPipelineOutput] thay vì một tuple thông thường hay không.
- use_resolution_binning (bool, tuỳ chọn, mặc định là True) — Nếu được đặt thành True, chiều cao và chiều rộng yêu cầu sẽ được ánh xạ trước tiên đến các độ phân giải gần nhất bằng cách sử dụng các bin tỷ lệ khung hình được định nghĩa trước. Sau khi các latents được tạo ra được giải mã thành hình ảnh, chúng sẽ được thay đổi kích thước trở lại độ phân giải yêu cầu. Hữu ích để tạo các hình ảnh không vuông ở các độ phân giải tối ưu.
- callback_on_step_end (Callable, optional) — Một hàm được gọi vào cuối mỗi bước khử nhiễu trong quá trình suy luận. Hàm này được gọi với các đối số sau: callback_on_step_end(self, step, timestep, callback_kwargs). callback_kwargs sẽ bao gồm một danh sách tất cả các tensor như được chỉ định bởi callback_on_step_end_tensor_inputs.
- callback_on_step_end_tensor_inputs (list, optional) — Danh sách các đầu vào tensor cho hàm callback_on_step_end. Các tensor được chỉ định trong danh sách sẽ được truyền dưới dạng tham số callback_kwargs. Bạn chỉ có thể bao gồm các tensor được liệt kê trong thuộc tính ._callback_tensor_inputs.
Trả về
[~pipelines.prx.PRXPipelineOutput] hoặc *tuple`
[~pipelines.prx.PRXPipelineOutput] nếu return_dict là True, ngược lại là một *tuple. Khi trả về một tuple, phần tử đầu tiên là một danh sách chứa các hình ảnh đã tạo.
Hàm được gọi khi gọi pipeline cho việc sinh.
Ví dụ:
>>> 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 )
Kiểm tra xem tất cả các đầu vào có đúng định dạng hay không.
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 )
Mã hóa prompt văn bản bằng trình mã hóa và bộ phân từ văn bản tiêu chuẩn, hoặc sử dụng các embedding đã được tính toán trước.
get_default_resolution
< source >( ) → int
Trả về
int
Kích thước mẫu (chiều cao/chiều rộng) mặc định được sử dụng cho quá trình tạo.
Xác định độ phân giải mặc định dựa trên VAE và cấu hình đã tải.
prepare_latents
< source >( batch_size: intnum_channels_latents: intheight: intwidth: intdtype: dtypedevice: devicegenerator: typing.Optional[torch.Generator] = Nonelatents: typing.Optional[torch.Tensor] = None )
Chuẩn bị các latent ban đầu cho quá trình khuếch tán.
PRXPipelineOutput
lớp diffusers.pipelines.prx.PRXPipelineOutput
< source >( images: list[PIL.Image.Image] | numpy.ndarray )
Lớp đầu ra cho các pipeline PRX.