Diffusers documentation

PixArt-Σ

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.43.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

PixArt-Σ

PixArt-Σ: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation の著者は、Junsong Chen、Jincheng Yu、Chongjian Ge、Lewei Yao、Enze Xie、Yue Wu、Zhongdao Wang、James Kwok、Ping Luo、Huchuan Lu、Zhenguo Li です。

論文の抄録は以下の通りです:

本論文では、4K解像度の画像を直接生成できるDiffusion Transformerモデル(DiT)であるPixArt-Σを提案します。PixArt-Σは、先行モデルであるPixArt-αから大きな進歩を遂げており、はるかに高い忠実度と、テキストプロンプトとの改善された整合性を提供します。PixArt-Σの重要な特徴の一つは、その学習効率です。PixArt-αの基礎的な事前学習を活用し、より高品質なデータを取り込むことで、「弱い」ベースラインから「強い」モデルへと進化させるプロセスを「weak-to-strong training」と呼んでいます。PixArt-Σの進歩は二つあります:(1) 高品質な学習データ:PixArt-Σは、より正確で詳細な画像キャプションと組み合わせた、より高品質な画像データを取り入れています。(2) 効率的なトークン圧縮:DiTフレームワーク内で、キーとバリューの両方を圧縮する新しいアテンションモジュールを提案し、効率を大幅に向上させ、超高分解能画像の生成を可能にしています。これらの改善により、PixArt-Σは、SDXL(2.6Bパラメータ)やSD Cascade(5.1Bパラメータ)などの既存のテキストから画像への拡散モデルよりもはるかに小さいモデルサイズ(0.6Bパラメータ)で、優れた画像品質とユーザープロンプトへの忠実性を達成します。さらに、PixArt-Σの4K画像生成能力は、高解像度のポスターや壁紙の作成を支援し、映画やゲームなどの業界における高品質なビジュアルコンテンツの制作を効率的に後押しします。

元のコードベースはPixArt-alpha/PixArt-sigma、利用可能なすべてのチェックポイントはPixArt-alphaで確認できます。

このパイプラインに関するいくつかの注記:

  • デノイジングにはUNetではなくTransformerバックボーンを使用します。そのため、DiTと同様のアーキテクチャを持っています。
  • T5から計算されたテキスト条件を用いて学習されています。この点により、複雑なテキストプロンプトや細かなディテールを含む指示をより正確に追従する能力が向上しています。
  • 異なるアスペクト比で高解像度の画像を生成する点に優れています。最適な結果を得るために、著者はいくつかのサイズ範囲を推奨しており、これらはこちらで確認できます。
  • 執筆時点において、PixArt-α、Stable Diffusion XL、Playground V2.0、DALL-E 3などの最先端のテキストから画像を生成するシステムと同等の品質を誇りながら、それらよりも効率的です。
  • 2048pxや4Kなど、超高分解能の画像を生成する能力を示しています。
  • テキストから画像を生成するモデルは、VAEやデータセットなどのいくつかの改善を通じて、弱いモデルからより強力なモデルへと成長できることを示しています。

スケジューラーのトレードオフ、すなわち速度と品質のバランスをどのように探求するかを学ぶには、ガイドを必ずご確認ください。また、同じコンポーネントを複数のパイプラインに効率的に読み込む方法については、パイプライン間でコンポーネントを再利用するセクションを参照してください。

PixArtSigmaPipelineで生成した画像をSDXL refinerモデルに渡すことで、生成品質をさらに向上させることができます。

8GB未満のGPU VRAMでの推論

8GB未満のGPU VRAMでPixArtSigmaPipelineを実行するには、テキストエンコーダーを8ビット精度で読み込みます。完全な例を順を追って見ていきましょう。

まず、bitsandbytesライブラリをインストールします:

pip install -U bitsandbytes

次に、テキストエンコーダーを8ビットで読み込みます:

from transformers import T5EncoderModel
from diffusers import PixArtSigmaPipeline
import torch

text_encoder = T5EncoderModel.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    subfolder="text_encoder",
    load_in_8bit=True,
    device_map="auto",
)
pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    text_encoder=text_encoder,
    transformer=None,
    device_map="balanced"
)

次に、pipe を使用してプロンプトをエンコードします:

with torch.no_grad():
    prompt = "cute cat"
    prompt_embeds, prompt_attention_mask, negative_embeds, negative_prompt_attention_mask = pipe.encode_prompt(prompt)

テキスト埋め込みが計算されたため、text_encoder と pipe をメモリから削除し、GPU VRAM を解放します:

import gc

def flush():
    gc.collect()
    torch.cuda.empty_cache()

del text_encoder
del pipe
flush()

次に、プロンプトの埋め込みを入力として、潜在変数(latents)を計算します:

pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    text_encoder=None,
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"

latents = pipe(
    negative_prompt=None,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    prompt_attention_mask=prompt_attention_mask,
    negative_prompt_attention_mask=negative_prompt_attention_mask,
    num_images_per_prompt=1,
    output_type="latent",
).images

del pipe.transformer
flush()

pipeを初期化する際、text_encoderをNoneに設定していることに注意してください。これにより、text_encoderは読み込まれません。

…

with torch.no_grad():
    image = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]
image = pipe.image_processor.postprocess(image, output_type="pil")[0]
image.save("cat.png")

使用していないコンポーネントを削除し、GPU VRAMを解放することで、PixArtSigmaPipelineを8GB未満のGPU VRAMで実行できるようになります。

メモリ使用量のレポートを取得するには、このスクリプトを実行してください。

8ビットで計算されたテキスト埋め込みは、精度の低下による表現空間の情報損失のため、生成される画像の品質に影響を与える可能性があります。8ビットを使用した場合と使用しない場合の出力を比較することをお勧めします。

text_encoderを読み込む際、load_in_8bitをTrueに設定しました。さらにメモリ要件を7GB未満まで下げたい場合は、load_in_4bitを指定することもできます。

PixArtSigmaPipeline

class diffusers.PixArtSigmaPipeline

< >

( tokenizer: T5Tokenizertext_encoder: T5EncoderModelvae: AutoencoderKLtransformer: PixArtTransformer2DModelscheduler: KarrasDiffusionSchedulers )

パラメータ

  • vae (AutoencoderKL) — 画像を潜在表現にエンコードし、潜在表現からデコードするための変分オートエンコーダ (VAE) モデル。
  • text_encoder (T5EncoderModel) — 凍結されたテキストエンコーダー。PixArt-Alpha は T5、具体的には t5-v1_1-xxl バリアントを使用します。
  • tokenizer (T5Tokenizer) — クラス T5Tokenizer のトークナイザー。
  • transformer (PixArtTransformer2DModel) — テキスト条件付きのPixArtTransformer2DModelで、エンコードされた画像の潜在変数(latents)のノイズ除去を行います。設定では当初Transformer2DModelとして公開されていましたが、この不一致は無視して問題ありません。
  • scheduler (SchedulerMixin) — transformer と組み合わせて、エンコードされた画像の潜在変数(latents)のノイズ除去を行うために使用するスケジューラ。

PixArt-Sigma を使用したテキストから画像を生成する Pipeline。

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

__call__

< >

( prompt: str | list[str] = Nonenegative_prompt: str = ''num_inference_steps: int = 20timesteps: list[int] = Nonesigmas: list[float] = Noneguidance_scale: float = 4.5num_images_per_prompt: int | None = 1height: int | None = Nonewidth: int | None = Noneeta: float = 0.0generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Nonenegative_prompt_embeds: typing.Optional[torch.Tensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.Tensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Truecallback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = Nonecallback_steps: int = 1clean_caption: bool = Trueuse_resolution_binning: bool = Truemax_sequence_length: int = 300**kwargs ) → ImagePipelineOutput or tuple

パラメータ

  • prompt (str または list[str]、省略可) — 画像生成をガイドするためのプロンプトです。未定義の場合、代わりに prompt_embeds を渡す必要があります。
  • negative_prompt (str または list[str]、省略可) — 画像生成をガイドしないプロンプトです。未定義の場合、代わりに negative_prompt_embeds を渡す必要があります。ガイダンスを使用していない場合は無視されます(guidance_scale が 1 未満の場合は無視されます)。
  • num_inference_steps (int、任意、デフォルト値は 100) — デノイジングのステップ数です。デノイジングのステップ数を増やすと、推論速度が遅くなる代わりに、通常はより高品質な画像が得られます。
  • timesteps (list[int], 任意) — timesteps 引数をサポートするスケジューラーの set_timesteps メソッドで、デノイジングプロセスに使用するカスタムタイムステップ。未定義の場合、num_inference_steps が渡された際のデフォルトの動作が使用されます。降順で指定する必要があります。
  • sigmas (list[float]、任意) — set_timesteps メソッドで sigmas 引数を受け付けるスケジューラーにおいて、デノイジングプロセスに使用するカスタムのシグマです。定義されていない場合、num_inference_steps が渡された際のデフォルトの動作が使用されます。
  • guidance_scale (float、任意、デフォルト値は 4.5) — Classifier-Free Diffusion Guidance で定義されているガイダンススケールです。guidance_scale は、Imagen Paper の数式 2 の w として定義されています。guidance_scale > 1 を設定することで、ガイダンススケールが有効になります。ガイダンススケールを高くすると、prompt のテキストと密接に関連する画像の生成が促されますが、通常は画像の品質が低下する傾向があります。
  • num_images_per_prompt (int、任意、デフォルトは 1) — プロンプトごとに生成する画像の数。
  • height (int、任意、デフォルトは self.unet.config.sample_size) — 生成される画像のピクセル単位の縦幅。
  • width (int、任意、デフォルトは self.unet.config.sample_size)— 生成される画像の幅(ピクセル単位)。
  • eta (float、省略可、デフォルト値は 0.0) — DDIM 論文におけるパラメータ eta (η) に相当します: https://huggingface.co/papers/2010.02502。 schedulers.DDIMScheduler にのみ適用され、それ以外の場合は無視されます。
  • generator (torch.Generator または list[torch.Generator]、省略可) — 生成を決定論的にするためのtorch ジェネレーター1つまたはリスト。
  • latents (torch.Tensor, 任意) — 画像生成の入力として使用する、ガウス分布からサンプリングされた事前生成されたノイズ付き潜在変数。異なるプロンプトで同じ生成を調整するために使用できます。指定されない場合、指定されたランダムgeneratorを使用してサンプリングを行い、潜在変数テンソルが生成されます。
  • prompt_embeds (torch.Tensor, 任意) — 事前に生成されたテキスト埋め込み。テキスト入力の調整、例えばプロンプトの重み付けなどを容易に行うために使用できます。指定されない場合、prompt 入力引数からテキスト埋め込みが生成されます。
  • prompt_attention_mask (torch.Tensor, 任意) — テキスト埋め込み用の事前生成されたアテンションマスク。
  • negative_prompt_embeds (torch.Tensor, 任意) — 事前に生成されたネガティブテキストの埋め込みです。PixArt-Sigma の場合、このネガティブプロンプトは "" である必要があります。指定されない場合、negative_prompt 入力引数から negative_prompt_embeds が生成されます。
  • negative_prompt_attention_mask (torch.Tensor, 任意) — ネガティブなテキスト埋め込み用の事前生成されたアテンションマスクです。
  • output_type (str、任意、デフォルトは "pil") — 画像生成の出力形式です。PIL: PIL.Image.Image または np.array から選択します。
  • return_dict (bool、省略可能、デフォルトは True) — ~pipelines.stable_diffusion.IFPipelineOutput を返すか、プレーンなタプルを返すかを指定します。
  • callback (Callable, 省略可) — 推論中に callback_steps ステップごとに呼び出される関数です。この関数は、以下の引数とともに呼び出されます: callback(step: int, timestep: int, latents: torch.Tensor)。
  • callback_steps (int、任意、デフォルトは 1) — callback 関数が呼び出される頻度です。指定しない場合、callback はすべてのステップで呼び出されます。
  • clean_caption (bool、任意、デフォルトは True) — 埋め込みを作成する前にキャプションをクリーンアップするかどうか。beautifulsoup4 と ftfy のインストールが必要です。依存関係がインストールされていない場合、埋め込みは生のプロンプトから作成されます。
  • use_resolution_binning (bool、デフォルトは True) — True に設定すると、要求された高さおよび幅は、ASPECT_RATIO_1024_BIN を使用して最も近い解像度にマッピングされます。生成された潜在変数が画像にデコードされた後、要求された解像도에リサイズされます。正方形でない画像の生成に有用です。
  • max_sequence_length (int、デフォルト値は 300)— prompt に使用する最大シーケンス長。

戻り値

ImagePipelineOutput または tuple

return_dict が True の場合、ImagePipelineOutput が返され、それ以外の場合は tuple が返されます。このとき、最初の要素は生成された画像のリストです。

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

例:

>>> import torch
>>> from diffusers import PixArtSigmaPipeline

>>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
>>> pipe = PixArtSigmaPipeline.from_pretrained(
...     "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
... )
>>> # Enable memory optimizations.
>>> # pipe.enable_model_cpu_offload()

>>> prompt = "A small cactus with a happy face in the Sahara desert."
>>> image = pipe(prompt).images[0]

encode_prompt

< >

( prompt: str | list[str]do_classifier_free_guidance: bool = Truenegative_prompt: str = ''num_images_per_prompt: int = 1device: typing.Optional[torch.device] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Nonenegative_prompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Nonenegative_prompt_attention_mask: typing.Optional[torch.Tensor] = Noneclean_caption: bool = Falsemax_sequence_length: int = 300**kwargs )

パラメータ

  • prompt (str または list[str]、省略可能) — エンコードするプロンプト
  • negative_prompt (str または list[str]、省略可) — 画像生成を導かないためのプロンプトです。未定義の場合、代わりに negative_prompt_embeds を渡す必要があります。ガイダンスを使用していない場合は無視されます(つまり、guidance_scale が 1 未満の場合は無視されます)。PixArt-Alpha の場合は、"" に設定してください。
  • do_classifier_free_guidance (bool、任意、デフォルトは True) — classifier free guidance を使用するかどうか
  • num_images_per_prompt (int、任意、デフォルトは 1) — プロンプトごとに生成する画像の数
  • device — (torch.device, 任意): 生成される埋め込みを配置する torch デバイス
  • prompt_embeds (torch.Tensor, 任意) — 事前に生成されたテキスト埋め込み。テキスト入力の調整、例えばプロンプトの重み付けなどを容易に行うために使用できます。指定されない場合、prompt 入力引数からテキスト埋め込みが生成されます。
  • negative_prompt_embeds (torch.Tensor, 任意) — 事前に生成されたネガティブなテキスト埋め込み。PixArt-Alpha の場合は、"" 文字列の埋め込みである必要があります。
  • clean_caption (bool、デフォルトは False) — True の場合、エンコード前に提供されたキャプションを前処理してクリーニングします。
  • max_sequence_length (int、デフォルト値は 300) — プロンプトに使用する最大シーケンス長。

プロンプトをテキストエンコーダーの隠れ状態にエンコードします。

Update on GitHub