Diffusers documentation
PRX
PRX
PRX генерирует изображения высокого качества по тексту, используя упрощённую архитектуру MMDIT, в которой токены текста не обновляются через блоки трансформера. Для эффективного сэмплирования применяется flow matching с дискретным планированием, а для многоязычного кодирования текста используется модель T5Gemma-2B-2B-UL2 от Google. Трансформер с ~1.3B параметров обеспечивает быстрое инференсирование без потери качества. Вы можете выбрать между Flux VAE (сжатие 8x, 16 латентных каналов) для баланса качества и скорости или DC-AE (сжатие 32x, 32 латентных канала) для латентного сжатия и более быстрой обработки.
Доступные модели
PRX предлагает несколько вариантов с различными конфигурациями VAE, каждый из которых оптимизирован для конкретных разрешений. Базовые модели лучше всего работают с детализированными промптами, позволяя улавливать сложные композиции и тонкие детали. Модели, дообученные на наборе данных Alchemist, повышают эстетическое качество, особенно при использовании более простых промптов.
| Модель | Разрешение | Доработанные | Дистиллированный | Описание | Рекомендуемые промпты | Рекомендуемые параметры | Рекомендуемый dtype |
|---|---|---|---|---|---|---|---|
Photoroom/prx-256-t2i | 256 | Нет | Нет | Базовая модель, предварительно обученная при 256 с VAE Flux | Лучше всего работает с подробными запросами на естественном языке | 28 шагов, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-256-t2i-sft | 512 | Да | Нет | Доработан на наборе данных Alchemist dataset с использованием Flux VAE | Способен обрабатывать менее детализированные запросы | 28 шагов, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i | 512 | Нет | Нет | Базовая модель, предварительно обученная при разрешении 512 с VAE Flux | Лучше всего работает с подробными запросами на естественном языке | 28 шагов, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-sft | 512 | Да | Нет | Доработан на наборе данных Alchemist dataset с использованием Flux VAE | Способен обрабатывать менее детализированные запросы на естественном языке | 28 шагов, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-sft-distilled | 512 | Да | Да | 8-шаговая дистиллированная модель из Photoroom/prx-512-t2i-sft | Способен обрабатывать менее детализированные запросы на естественном языке | 8 шагов, cfg=1.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae | 512 | Нет | Нет | Базовая модель, предварительно обученная при разрешении 512 с использованием Deep Compression Autoencoder (DC-AE) | Лучше всего работает с подробными запросами на естественном языке | 28 шагов, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae-sft | 512 | Да | Нет | Доработан на наборе данных Alchemist dataset с использованием Deep Compression Autoencoder (DC-AE) | Способен обрабатывать менее детализированные запросы на естественном языке | 28 шагов, cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-dc-ae-sft-distilled | 512 | Да | Да | 8-шаговая дистиллированная модель из Photoroom/prx-512-t2i-dc-ae-sft-distilled | Способен обрабатывать менее детализированные запросы на естественном языке | 8 шагов, cfg=1.0 | torch.bfloat16 |
См. это собрание для получения дополнительной информации.
Загрузка конвейера
Загрузите конвейер с помощью 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")Ручная загрузка компонентов
Загружайте компоненты по отдельности, чтобы настроить конвейер, например, для использования квантованных моделей.
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"Оптимизация памяти
Для сред с ограниченной памятью:
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
< 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 )
Параметры
- transformer (
PRXTransformer2DModel) — Модель PRX transformer для удаления шума из закодированных латентов изображения. - scheduler (FlowMatchEulerDiscreteScheduler) —
планировщик, используемый в связке с
transformerдля удаления шума из закодированных латентов изображения. - text_encoder (
T5GemmaEncoder) — Модель текстового кодировщика для кодирования запросов. - tokenizer ([
T5TokenizerFastилиGemmaTokenizerFast]) — Токенизатор для текстового энкодера. - vae (AutoencoderKL or AutoencoderDC) — Модель вариационного автоэнкодера (VAE) для кодирования изображений в латентные представления и декодирования из них. Поддерживает как AutoencoderKL (сжатие 8x), так и AutoencoderDC (сжатие 32x).
Пайплайн для генерации изображений по тексту с использованием PRX Transformer.
Эта модель наследуется от DiffusionPipeline. См. документацию родительского класса для общих методов, которые библиотека реализует для всех пайплайнов (например, загрузка или сохранение, запуск на конкретном устройстве и т. д.)
__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`
Параметры
- prompt (str или list[str], необязательный) — Промпт или промпты для управления генерацией изображения. Если не задан, необходимо передать prompt_embeds вместо него.
- negative_prompt (str, необязательный, по умолчанию "") — Промпт или промпты, которые не должны направлять генерацию изображения. Игнорируется, если не используется направляющее масштабирование (то есть игнорируется, если guidance_scale меньше 1).
- height (int, optional, по умолчанию self.transformer.config.sample_size * self.vae_scale_factor) — Высота генерируемого изображения в пикселях.
- width (int, необязательный, по умолчанию self.transformer.config.sample_size * self.vae_scale_factor) — Ширина генерируемого изображения в пикселях.
- num_inference_steps (int, optional, по умолчанию 28) — Количество шагов денойзинга. Как правило, большее число шагов денойзинга приводит к изображению более высокого качества, но за счёт более медленного вывода.
- timesteps (list[int], optional) — Пользовательские шаги для процесса удаления шума с планировщиками, поддерживающими аргумент timesteps в их методе set_timesteps. Если не задано, будет использоваться поведение по умолчанию при передаче num_inference_steps. Должны быть в порядке убывания.
- guidance_scale (float, optional, по умолчанию 4.0) — Масштаб руководства, определенный в Classifier-Free Diffusion Guidance. guidance_scale определяется как w из уравнения 2 статьи Imagen Paper. Масштаб руководства включается при установке guidance_scale > 1. Более высокое значение масштаба руководства способствует генерации изображений, тесно связанных с текстом prompt, обычно за счет снижения качества изображения.
- num_images_per_prompt (int, optional, по умолчанию 1) — Количество изображений, генерируемых для каждого промпта.
- generator (torch.Generator или list[torch.Generator], необязательный) — Один или список генераторов torch для обеспечения детерминированности генерации.
- latents (torch.Tensor, optional) — Предварительно сгенерированные зашумленные латенты, выбранные из гауссова распределения, которые используются в качестве входных данных для генерации изображений. Позволяет корректировать одну и ту же генерацию с помощью разных промптов. Если не заданы, тензор латентов будет сгенерирован путём выборки с использованием переданного случайного generator.
- prompt_embeds (torch.FloatTensor, optional) — Предварительно сгенерированные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, например, для взвешивания промпта. Если не предоставлены, текстовые эмбеддинги будут сгенерированы из входного аргумента prompt.
- negative_prompt_embeds (torch.FloatTensor, optional) — Предварительно сгенерированные отрицательные текстовые эмбеддинги. Можно использовать для простой настройки текстовых входных данных, например для взвешивания промпта. Если не указано и guidance_scale > 1, отрицательные эмбеддинги будут сгенерированы из пустой строки.
- prompt_attention_mask (torch.BoolTensor, optional) — Предварительно сгенерированная маска внимания для prompt_embeds. Если не задана, маска внимания будет сгенерирована из аргумента prompt.
- negative_prompt_attention_mask (torch.BoolTensor, optional) — Предварительно сгенерированная маска внимания для negative_prompt_embeds. Если не задана и guidance_scale > 1, маска внимания будет сгенерирована из пустой строки.
- output_type (str, optional, по умолчанию “pil”) — Формат вывода для генерируемого изображения. Выберите один из вариантов: PIL: PIL.Image.Image или np.array.
- tokenizer_max_length (int, optional) —
Переопределяет максимальное количество токенов, используемых при токенизации промпта. По умолчанию используется собственное значение
model_max_lengthтокенизатора, если параметр не задан. - skip_text_cleaning (bool, optional, по умолчанию False) — Если True, используется только лёгкая очистка промпта (исправление кодировки + снятие HTML-экранирования) вместо полного конвейера очистки DeepFloyd.
- return_dict (bool, optional, по умолчанию True) — Определяет, следует ли возвращать [~pipelines.prx.PRXPipelineOutput] вместо обычного кортежа.
- use_resolution_binning (bool, optional, по умолчанию True) — Если установлено в True, запрашиваемые высота и ширина сначала сопоставляются с ближайшими разрешениями с использованием предопределённых бина соотношений сторон. После декодирования полученных латентов в изображения они масштабируются обратно до запрошенного разрешения. Полезно для генерации несимметричных изображений с оптимальными разрешениями.
- callback_on_step_end (Callable, optional) — Функция, вызываемая в конце каждого шага денойзинга во время инференса. Функция вызывается со следующими аргументами: callback_on_step_end(self, step, timestep, callback_kwargs). callback_kwargs будет содержать список всех тензоров, как указано в callback_on_step_end_tensor_inputs.
- callback_on_step_end_tensor_inputs (list, optional) — Список тензорных входных данных для функции callback_on_step_end. Тензоры, указанные в списке, будут переданы в качестве аргумента callback_kwargs. Вы сможете включать только те тензоры, которые перечислены в атрибуте ._callback_tensor_inputs.
Возвращает
[~pipelines.prx.PRXPipelineOutput] или *tuple`
[~pipelines.prx.PRXPipelineOutput], если return_dict равно True, иначе — *кортеж. При возврате кортежа первым элементом является список с сгенерированными изображениями.
Функция, вызываемая при обращении к конвейеру для генерации.
Примеры:
>>> 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 )
Убедитесь, что все входные данные имеют правильный формат.
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 )
Кодируйте текстовый промпт с помощью стандартного текстового энкодера и токенизатора или используйте предварительно вычисленные эмбеддинги.
get_default_resolution
< source >( ) → int
Возвращает
int
Размер выборки (высота/ширина) по умолчанию, используемый для генерации.
Определите разрешение по умолчанию на основе загруженной VAE и конфигурации.
prepare_latents
< source >( batch_size: intnum_channels_latents: intheight: intwidth: intdtype: dtypedevice: devicegenerator: typing.Optional[torch.Generator] = Nonelatents: typing.Optional[torch.Tensor] = None )
Подготовьте начальные латенты для процесса диффузии.
PRXPipelineOutput
класс diffusers.pipelines.prx.PRXPipelineOutput
< source >( images: list[PIL.Image.Image] | numpy.ndarray )
Класс вывода для конвейеров PRX.