Diffusers documentation
PRX
PRX
PRXは、テキストトークンがTransformerブロックを通じて更新されない簡素化されたMMDITアーキテクチャを用いて、テキストから高品質な画像を生成します。効率的なサンプリングのために、離散スケジューリングを備えたフローマッチングを採用し、多言語テキストエンコーディングにはGoogleのT5Gemma-2B-2B-UL2モデルを使用しています。約1.3BパラメータのTransformerにより、品質を犠牲にすることなく高速な推論を実現します。バランスの取れた品質と速度を求める場合はFlux VAE（8倍圧縮、16の潜在チャネル）、潜在圧縮とより高速な処理を求める場合はDC-AE（32倍圧縮、32の潜在チャネル）を選択できます。
利用可能なモデル
PRXには、異なるVAE設定を持つ複数のバリアントがあり、それぞれが特定の解像度に最適化されています。ベースモデルは詳細なプロンプトで優れた性能を発揮し、複雑な構図や微細なディテールを捉えます。Alchemist datasetで学習されたファインチューニング済みモデルは、特にシンプルなプロンプトにおいて、美的品質を向上させます。
| モデル | 解決 | ファインチューニング済み | 蒸留済み | 説明 | 推奨プロンプト | 推奨パラメータ | 推奨されるdtype |
|---|---|---|---|---|---|---|---|
Photoroom/prx-256-t2i | 256 | いいえ | いいえ | Flux VAEで256に事前学習されたベースモデル | 自然言語による詳細なプロンプトで最も効果的に動作します | 28ステップ、cfg=5.0 | torch.bfloat16 |
Photoroom/prx-256-t2i-sft | 512 | はい | いいえ | Flux VAE を使用して、Alchemist dataset でファインチューニング | 詳細度が低いプロンプトにも対応できます | 28ステップ、cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i | 512 | いいえ | いいえ | Flux VAEで512に事前学習されたベースモデル | 自然言語による詳細なプロンプトで最も効果的に動作します | 28ステップ、cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-sft | 512 | はい | いいえ | Flux VAE を使用して、Alchemist dataset でファインチューニング | 自然言語による詳細度の低いプロンプトを処理できます | 28ステップ、cfg=5.0 | torch.bfloat16 |
Photoroom/prx-512-t2i-sft-distilled | 512 | はい | はい | Photoroom/prx-512-t2i-sftからの8ステップ蒸留モデル | 自然言語による詳細度の低いプロンプトを処理できます | 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 | はい | はい | Photoroom/prx-512-t2i-dc-ae-sft-distilledからの8ステップ蒸留モデル | 自然言語による詳細度の低いプロンプトを処理できます | 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
class 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) — エンコードされた画像の潜在変数（latents）をデノイズするための PRX transformer モデル。 - scheduler (FlowMatchEulerDiscreteScheduler) —
transformerと組み合わせて、エンコードされた画像の潜在変数からノイズを除去するために使用するスケジューラ。 - text_encoder (
T5GemmaEncoder) — プロンプトをエンコードするためのテキストエンコーダーモデルです。 - tokenizer ([
T5TokenizerFastまたはGemmaTokenizerFast]) — テキストエンコーダー用のトークナイザー。 - vae (AutoencoderKL または AutoencoderDC) — 画像を潜在表現にエンコードし、潜在表現からデコードするための Variational Auto-Encoder (VAE) モデルです。 AutoencoderKL (8倍圧縮) と AutoencoderDC (32倍圧縮) の両方をサポートしています。
PRX Transformer を使用したテキストから画像への生成を行う Pipeline。
このモデルは 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]、optional) — 画像生成をガイドするためのプロンプトです。未定義の場合、代わりに prompt_embeds を渡す必要があります。
- negative_prompt (str、optional、デフォルトは "") — 画像生成をガイドしないプロンプトです。ガイダンスを使用していない場合は無視されます（guidance_scale が 1 未満の場合は無視されます）。
- height (int、任意、デフォルトは 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]、任意) — timesteps 引数をサポートするスケジューラーの set_timesteps メソッドで、デノイジングプロセスに使用するカスタムタイムステップ。定義されていない場合、num_inference_steps が渡された際のデフォルトの動作が使用されます。降順で指定する必要があります。
- guidance_scale (float、optional、デフォルト値は 4.0) — Classifier-Free Diffusion Guidance で定義されているガイダンススケールです。guidance_scale は、Imagen Paper の数式 2 の w として定義されています。guidance_scale > 1 を設定することでガイダンススケールが有効になります。ガイダンススケールを高くすると、prompt のテキストと強く関連する画像の生成が促されますが、通常は画像の品質が低下する傾向があります。
- num_images_per_prompt (int、任意、デフォルト値は 1) — プロンプトごとに生成する画像の数。
- generator (torch.Generator または list[torch.Generator]、省略可) — 生成を決定論的にするための torch ジェネレーター 1 つまたはリスト。
- latents (torch.Tensor, optional) — 画像生成の入力として使用する、ガウス分布からサンプリングされた事前生成されたノイジーな潜在変数です。異なるプロンプトで同じ生成を調整するために使用できます。指定されない場合、指定されたランダムgeneratorを使用してサンプリングを行い、潜在変数テンソルが生成されます。
- prompt_embeds (torch.FloatTensor、任意) — 事前に生成されたテキスト埋め込み。テキスト入力の調整、例えばプロンプトの重み付けなどを容易に行うことができます。指定されない場合、prompt 入力引数からテキスト埋め込みが生成されます。
- negative_prompt_embeds (torch.FloatTensor、任意) — 事前に生成されたネガティブなテキスト埋め込み。テキスト入力の調整、例えばプロンプトの重み付けなどを容易に行うために使用できます。指定されず、かつ guidance_scale > 1 の場合、空の文字列からネガティブな埋め込みが生成されます。
- prompt_attention_mask (torch.BoolTensor、任意) — prompt_embeds 用の事前生成されたアテンションマスクです。指定されない場合、prompt 入力引数からアテンションマスクが生成されます。
- negative_prompt_attention_mask (torch.BoolTensor、任意) — 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、任意、デフォルトは False) — True の場合、DeepFloyd の完全なクリーニングパイプラインの代わりに、軽量のプロンプトクリーニング（エンコーディングの修正と HTML のアンエスケープ）のみを使用します。
- 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 の場合、それ以外の場合は *tuple を返します。tuple を返す場合、最初の要素は生成された画像のリストです。
生成用のパイプラインを呼び出したときに実行される関数。
例:
>>> 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 )
標準的なテキストエンコーダーとトークナイザーを使用してテキストプロンプトをエンコードするか、事前計算された埋め込みを使用します。
読み込まれた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 パイプライン用の出力クラスです。