#
# Apache License, Version 2.0（「ライセンス」）の下でライセンスされています。
# このファイルは、ライセンスに準拠した場合にのみ使用できます。
# ライセンスのコピーは以下から入手できます
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 適用される法律で要求される場合、または書面による合意がある場合を除き、ソフトウェア
# ライセンスに基づき配布される本ソフトウェアは、「現状のまま」
# 明示的または黙示的を問わず、いかなる種類の保証や条件も付随しません。
# 権限を管理する特定の言語については、ライセンスを参照してください。
# ライセンスの下での制限。-->

# PRX

PRXは、テキストトークンがTransformerブロックを通じて更新されない簡素化されたMMDITアーキテクチャを用いて、テキストから高品質な画像を生成します。効率的なサンプリングのために、離散スケジューリングを用いたフローマッチングを採用し、多言語テキストエンコーディングにはGoogleのT5Gemma-2B-2B-UL2モデルを使用しています。約1.3BパラメータのTransformerにより、品質を犠牲にすることなく高速な推論を実現します。バランスの取れた品質と速度を求める場合はFlux VAE（8倍圧縮、16の潜在チャネル）、潜在圧縮とより高速な処理を求める場合はDC-AE（32倍圧縮、32の潜在チャネル）を選択できます。

## 利用可能なモデル

PRXには、異なるVAE設定を持つ複数のバリアントがあり、それぞれが特定の解像度に最適化されています。ベースモデルは詳細なプロンプトで優れた性能を発揮し、複雑な構図や微細なディテールを捉えます。[Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist)で学習されたファインチューニング済みモデルは、特にシンプルなプロンプトにおいて、美的品質を向上させます。

| モデル | 解決 | ファインチューニング済み | 蒸留済み | 説明 | 推奨プロンプト | 推奨パラメータ | 推奨されるdtype |
|:-----:|:-----------------:|:----------:|:----------:|:----------:|:----------:|:----------:|:----------:|
| [`Photoroom/prx-256-t2i`](https://huggingface.co/Photoroom/prx-256-t2i)| 256 | いいえ | いいえ | Flux VAEで256に事前学習されたベースモデル|自然言語による詳細なプロンプトで最も効果的に動作します|28ステップ、cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-256-t2i-sft`](https://huggingface.co/Photoroom/prx-256-t2i-sft)| 512 | はい | いいえ | Flux VAE を使用して、[Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) でファインチューニング | 詳細度が低いプロンプトにも対応できます|28ステップ、cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i`](https://huggingface.co/Photoroom/prx-512-t2i)| 512 | いいえ | いいえ | Flux VAEで512に事前学習されたベースモデル |自然言語による詳細なプロンプトで最も効果的に動作します|28ステップ、cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft)| 512 | はい | いいえ | Flux VAE を使用して、[Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist) でファインチューニング | 自然言語による詳細度の低いプロンプトを処理できます|28ステップ、cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-sft-distilled)| 512 | はい | はい | [`Photoroom/prx-512-t2i-sft`](https://huggingface.co/Photoroom/prx-512-t2i-sft)からの8ステップ蒸留モデル | 自然言語による詳細度の低いプロンプトを処理できます|8ステップ、cfg=1.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae)| 512 | いいえ | いいえ | 512で事前学習されたベースモデル、[Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)|自然言語による詳細なプロンプトで最も効果的に動作します|28ステップ、cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae-sft`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft)| 512 | はい | いいえ | [Alchemist dataset](https://huggingface.co/datasets/yandex/alchemist)で、[Deep Compression Autoencoder (DC-AE)](https://hanlab.mit.edu/projects/dc-ae)を使用してファインチューニングしました | 自然言語による詳細度の低いプロンプトを処理できます|28ステップ、cfg=5.0| `torch.bfloat16` |
| [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled)| 512 | はい | はい | [`Photoroom/prx-512-t2i-dc-ae-sft-distilled`](https://huggingface.co/Photoroom/prx-512-t2i-dc-ae-sft-distilled)からの8ステップ蒸留モデル | 自然言語による詳細度の低いプロンプトにも対応可能 | 8ステップ、cfg=1.0 | `torch.bfloat16` |s

詳細については、[こちら](https://huggingface.co/collections/Photoroom/prx-models-68e66254c202ebfab99ad38e)のコレクションを参照してください。

## パイプラインの読み込み

[from_pretrained()](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained)でパイプラインを読み込みます。

```py
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")
```

### 手動コンポーネントの読み込み

パイプラインをカスタマイズするため（例えば、量子化モデルを使用する場合など）、コンポーネントを個別に読み込むことができます。

```py
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"
```

## メモリ最適化

メモリが制約された環境向け:

```py
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]]

#### diffusers.PRXPipeline[[diffusers.PRXPipeline]]

```python
diffusers.PRXPipeline(transformer: PRXTransformer2DModel, scheduler: FlowMatchEulerDiscreteScheduler, text_encoder: T5GemmaEncoder, tokenizer: transformers.models.t5.tokenization_t5.T5Tokenizer | transformers.models.gemma.tokenization_gemma.GemmaTokenizer | transformers.models.auto.tokenization_auto.AutoTokenizer, vae: diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL | diffusers.models.autoencoders.autoencoder_dc.AutoencoderDC | None = None, default_sample_size: int | None = 512)
```

[ソース](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/prx/pipeline_prx.py#L257)

**パラメータ:**

transformer (`PRXTransformer2DModel`) : エンコードされた画像の潜在変数をノイズ除去するためのPRX transformerモデル。

scheduler ([FlowMatchEulerDiscreteScheduler](/docs/diffusers/main/ja/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler)) : `transformer` と組み合わせて、エンコードされた画像の潜在変数をデノイズするために使用するスケジューラー。

text_encoder (`T5GemmaEncoder`) : プロンプトをエンコードするためのテキストエンコーダーモデル。

tokenizer ([`T5TokenizerFast` または `GemmaTokenizerFast`]) : テキストエンコーダー用のトークナイザー。

vae ([AutoencoderKL](/docs/diffusers/main/ja/api/models/autoencoderkl#diffusers.AutoencoderKL) または [AutoencoderDC](/docs/diffusers/main/ja/api/models/autoencoder_dc#diffusers.AutoencoderDC)) : 画像を潜在表現にエンコードし、潜在表現からデコードするための変分オートエンコーダ (VAE) モデルです。AutoencoderKL (8 倍圧縮) と AutoencoderDC (32 倍圧縮) の両方をサポートしています。

PRX Transformer を使用したテキストから画像への生成を行う Pipeline。

このモデルは [DiffusionPipeline](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline) を継承しています。ライブラリがすべてのパイプラインに対して実装している汎用的なメソッド（ダウンロードや保存、特定のデバイスでの実行など）については、スーパークラスのドキュメントを参照してください。

#### __call__[[diffusers.PRXPipeline.__call__]]

```python
__call__(prompt: str | list[str] = None, negative_prompt: str = '', height: int | None = None, width: int | None = None, num_inference_steps: int = 28, timesteps: list = None, guidance_scale: float = 4.0, num_images_per_prompt: int | None = 1, generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = None, latents: typing.Optional[torch.Tensor] = None, prompt_embeds: typing.Optional[torch.FloatTensor] = None, negative_prompt_embeds: typing.Optional[torch.FloatTensor] = None, prompt_attention_mask: typing.Optional[torch.BoolTensor] = None, negative_prompt_attention_mask: typing.Optional[torch.BoolTensor] = None, output_type: str | None = 'pil', return_dict: bool = True, use_resolution_binning: bool = True, callback_on_step_end: typing.Optional[typing.Callable[[int, int], NoneType]] = None, callback_on_step_end_tensor_inputs: list = ['latents'], tokenizer_max_length: int | None = None, skip_text_cleaning: bool = False)
```

[ソース](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/prx/pipeline_prx.py#L554)

**パラメータ:**

prompt (*str* または *list[str]*、*任意*) : 画像生成をガイドするためのプロンプトです。定義されていない場合は、代わりに *prompt_embeds* を渡す必要があります。

negative_prompt (*str*、*optional*、デフォルトは *""*) : 画像生成をガイドしないためのプロンプトまたはプロンプツ。ガイダンスを使用していない場合は無視されます（*guidance_scale* が *1* 未満の場合は無視されます）。

height (*int*、*optional*、デフォルトは self.transformer.config.sample_size * self.vae_scale_factor) : 生成される画像のピクセル単位の縦幅。

width (*int*、*optional*、デフォルトは self.transformer.config.sample_size * self.vae_scale_factor) : 生成される画像の幅（ピクセル単位）。

num_inference_steps (*int*、*optional*、デフォルト値は 28) : デノイジングのステップ数です。デノイジングのステップ数を増やすと、推論速度が遅くなる代わりに、通常はより高品質な画像が得られます。

timesteps (*list[int]*、*optional*) : *set_timesteps* メソッドで *timesteps* 引数をサポートするスケジューラーにおいて、デノイジングプロセスに使用するカスタムタイムステップ。*num_inference_steps* が渡された場合のデフォルトの動作が使用されます。降順で指定する必要があります。

guidance_scale (*float*、*optional*、デフォルト値は 4.0) : [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598) で定義されているガイダンススケールです。*guidance_scale* は、[Imagen Paper](https://huggingface.co/papers/2205.11487) の数式 2 の *w* として定義されています。*guidance_scale > 1* を設定することで、ガイダンススケールが有効になります。ガイダンススケールを高くすると、*prompt* のテキストと密接に関連する画像の生成が促されますが、通常は画像の品質が低下する傾向があります。

num_images_per_prompt (*int*、*optional*、デフォルトは 1) : プロンプトごとに生成する画像の数。

generator (*torch.Generator* または *list[torch.Generator]*、*任意*) : 生成を決定論的にするための [torch ジェネレーター](https://pytorch.org/docs/stable/generated/torch.Generator.html) 1 つまたはそのリスト。

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](https://pillow.readthedocs.io/en/stable/): *PIL.Image.Image* または *np.array* から選択します。

tokenizer_max_length (*int*, *optional*) : プロンプトのトークン化時に使用される最大トークン数を上書きします。設定されていない場合、デフォルトではトークナイザー自身の`model_max_length`が使用されます。

skip_text_cleaning (*bool*、*optional*、デフォルトは *False*) : *True* の場合、完全な DeepFloyd クリーニングパイプラインの代わりに、軽量なプロンプトクリーニング（エンコーディングの修正 + HTML のアンエスケープ）のみを使用します。

return_dict (*bool*、*optional*、デフォルトは *True*) : プレーンなタプルではなく、[*~pipelines.prx.PRXPipelineOutput*] を返すかどうか。

use_resolution_binning (*bool*、*optional*、デフォルトは *True*) : *True* に設定すると、要求された高さおよび幅は、事前に定義されたアスペクト比のビンを使用して、最も近い解像度にマッピングされます。生成された潜在変数が画像にデコードされた後、要求された解像도에リサイズされます。最適な解像度で正方形でない画像を生成する場合に有用です。

callback_on_step_end (*Callable*、*任意*) : 推論中の各デノイジングステップの終了時に呼び出される関数です。この関数は、以下の引数とともに呼び出されます: *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 を返す場合、最初の要素は生成された画像のリストです。

生成用のパイプラインを呼び出したときに実行される関数。

例:
```py
>>> 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[[diffusers.PRXPipeline.check_inputs]]

```python
check_inputs(prompt: str | list[str], height: int, width: int, guidance_scale: float, callback_on_step_end_tensor_inputs: list[str] | None = None, prompt_embeds: typing.Optional[torch.FloatTensor] = None, negative_prompt_embeds: typing.Optional[torch.FloatTensor] = None)
```

[ソース](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/prx/pipeline_prx.py#L500)

すべての入力が正しい形式であることを確認してください。

#### encode_prompt[[diffusers.PRXPipeline.encode_prompt]]

```python
encode_prompt(prompt: str | list[str], device: typing.Optional[torch.device] = None, do_classifier_free_guidance: bool = True, negative_prompt: str = '', num_images_per_prompt: int = 1, prompt_embeds: typing.Optional[torch.FloatTensor] = None, negative_prompt_embeds: typing.Optional[torch.FloatTensor] = None, prompt_attention_mask: typing.Optional[torch.BoolTensor] = None, negative_prompt_attention_mask: typing.Optional[torch.BoolTensor] = None, tokenizer_max_length: int | None = None, skip_text_cleaning: bool = False)
```

[ソース](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/prx/pipeline_prx.py#L376)

標準的なテキストエンコーダーとトークナイザーを使用してテキストプロンプトをエンコードするか、事前計算された埋め込みを使用します。

#### get_default_resolution[[diffusers.PRXPipeline.get_default_resolution]]

```python
get_default_resolution()
```

[ソース](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/prx/pipeline_prx.py#L340)

**戻り値:** `int`

生成に使用するデフォルトのサンプルサイズ（高さ/幅）。

読み込まれたVAEと設定に基づいて、デフォルトの解像度を決定します。

#### prepare_latents[[diffusers.PRXPipeline.prepare_latents]]

```python
prepare_latents(batch_size: int, num_channels_latents: int, height: int, width: int, dtype: dtype, device: device, generator: typing.Optional[torch.Generator] = None, latents: typing.Optional[torch.Tensor] = None)
```

[ソース](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/prx/pipeline_prx.py#L352)

拡散プロセス用の初期潜在変数を準備します。

## PRXPipelineOutput[[diffusers.pipelines.prx.PRXPipelineOutput]]

#### diffusers.pipelines.prx.PRXPipelineOutput[[diffusers.pipelines.prx.PRXPipelineOutput]]

```python
diffusers.pipelines.prx.PRXPipelineOutput(images: list[PIL.Image.Image] | numpy.ndarray)
```

[ソース](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/prx/pipeline_output.py#L24)

**パラメータ:**

images (`list[PIL.Image.Image]` または `np.ndarray`) : 長さ `batch_size` のデノイズ済み PIL イメージのリスト、または形状 `(batch_size, height, width, num_channels)` の numpy 配列。PIL イメージまたは numpy 配列は、拡散パイプラインのデノイズ済みイメージを表します。

PRX パイプライン用の出力クラスです。