Diffusers 文件
PixArt-α
並獲得增強的文件體驗
開始使用
PixArt-α
PixArt-α:快速訓練擴散Transformer以實現逼真文字到影像合成由Junsong Chen、Jincheng Yu、Chongjian Ge、Lewei Yao、Enze Xie、Yue Wu、Zhongdao Wang、James Kwok、Ping Luo、Huchuan Lu 和 Zhenguo Li 完成。
論文摘要如下:
最先進的文字到影像 (T2I) 模型需要巨大的訓練成本(例如,數百萬 GPU 小時),嚴重阻礙了 AIGC 社群的基礎創新,同時增加了二氧化碳排放。本文介紹了 PIXART-α,一個基於 Transformer 的 T2I 擴散模型,其影像生成質量與最先進的影像生成器(例如 Imagen、SDXL 甚至 Midjourney)具有競爭力,達到了接近商業應用的水平。此外,它支援高達 1024 畫素的高解析度影像合成,且訓練成本低,如圖 1 和圖 2 所示。為實現這一目標,提出了三個核心設計:(1)訓練策略分解:我們設計了三個不同的訓練步驟,分別最佳化畫素依賴性、文字影像對齊和影像美學質量;(2)高效 T2I Transformer:我們將交叉注意力模組整合到擴散 Transformer (DiT) 中,以注入文字條件並簡化計算密集型類別條件分支;(3)高資訊量資料:我們強調文字影像對中概念密度的重要性,並利用大型視覺語言模型自動標記密集偽字幕以輔助文字影像對齊學習。結果,PIXART-α 的訓練速度明顯超過現有的大規模 T2I 模型,例如 PIXART-α 僅需 Stable Diffusion v1.5 訓練時間的 10.8%(675 對比 6,250 A100 GPU 天),節省了近 30 萬美元(26,000 美元對比 320,000 美元),並減少了 90% 的二氧化碳排放。此外,與更大的 SOTA 模型 RAPHAEL 相比,我們的訓練成本僅為 1%。大量實驗表明,PIXART-α 在影像質量、藝術性和語義控制方面表現出色。我們希望 PIXART-α 能為 AIGC 社群和初創公司提供新思路,以加速從零開始構建自己的高質量、低成本生成模型。
您可以在 PixArt-alpha/PixArt-alpha 找到原始程式碼庫,在 PixArt-alpha 找到所有可用的檢查點。
關於此管道的一些注意事項
- 它使用 Transformer 主幹(而不是 UNet)進行去噪。因此,它的架構與 DiT 相似。
- 它使用從 T5 計算的文字條件進行訓練。這使得該管道更擅長遵循具有複雜細節的複雜文字提示。
- 它擅長生成不同縱橫比的高解析度影像。為了獲得最佳結果,作者推薦了一些尺寸範圍,可以在這裡找到。
- 它在影像生成質量方面可與當前最先進的文字到影像生成系統(截至撰寫本文時)如 Stable Diffusion XL、Imagen 和 DALL-E 2 媲美,同時效率更高。
使用低於 8GB GPU 視訊記憶體進行推理
透過以 8 位精度載入文字編碼器,使用低於 8GB GPU 視訊記憶體執行 PixArtAlphaPipeline。讓我們來看一個完整的例子。
首先,安裝 bitsandbytes 庫
pip install -U bitsandbytes
然後以 8 位載入文字編碼器
from transformers import T5EncoderModel
from diffusers import PixArtAlphaPipeline
import torch
text_encoder = T5EncoderModel.from_pretrained(
"PixArt-alpha/PixArt-XL-2-1024-MS",
subfolder="text_encoder",
load_in_8bit=True,
device_map="auto",
)
pipe = PixArtAlphaPipeline.from_pretrained(
"PixArt-alpha/PixArt-XL-2-1024-MS",
text_encoder=text_encoder,
transformer=None,
device_map="auto"
)
現在,使用 `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 視訊記憶體
import gc
def flush():
gc.collect()
torch.cuda.empty_cache()
del text_encoder
del pipe
flush()
然後以提示嵌入作為輸入計算潛在變數
pipe = PixArtAlphaPipeline.from_pretrained(
"PixArt-alpha/PixArt-XL-2-1024-MS",
text_encoder=None,
torch_dtype=torch.float16,
).to("cuda")
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`,這樣它就不會被載入。
潛在變數計算完成後,將其傳遞給 VAE 進行解碼,生成真實影像
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 視訊記憶體,您應該能夠以低於 8GB GPU 視訊記憶體執行 PixArtAlphaPipeline。
如果您需要記憶體使用情況報告,請執行此指令碼。
以 8 位計算文字嵌入可能會影響生成影像的質量,因為精度降低導致表示空間中的資訊丟失。建議比較 8 位和非 8 位下的輸出。
載入 `text_encoder` 時,您將其 `load_in_8bit` 設定為 `True`。您也可以指定 `load_in_4bit`,以將記憶體需求進一步降低到 7GB 以下。
PixArtAlphaPipeline
類 diffusers.PixArtAlphaPipeline
< 源 >( tokenizer: T5Tokenizer text_encoder: T5EncoderModel vae: AutoencoderKL transformer: PixArtTransformer2DModel scheduler: DPMSolverMultistepScheduler )
引數
- vae (AutoencoderKL) — 用於將影像編碼和解碼為潛在表示的變分自編碼器 (VAE) 模型。
- text_encoder (
T5EncoderModel
) — 凍結的文字編碼器。PixArt-Alpha 使用 T5,特別是 t5-v1_1-xxl 變體。 - 分詞器 (
T5Tokenizer
) — T5Tokenizer 類的分詞器。 - transformer (PixArtTransformer2DModel) — 一個文字條件
PixArtTransformer2DModel
,用於對編碼的影像潛在變數進行去噪。 - 排程器 (SchedulerMixin) — 用於與
transformer
結合對編碼影像潛在變數進行去噪的排程器。
用於使用 PixArt-Alpha 生成文字到影像的管道。
此模型繼承自 DiffusionPipeline。請檢視超類文件,瞭解該庫為所有管道實現的通用方法(例如下載或儲存、在特定裝置上執行等)。
__call__
< 源 >( prompt: typing.Union[str, typing.List[str]] = None negative_prompt: str = '' num_inference_steps: int = 20 timesteps: typing.List[int] = None sigmas: typing.List[float] = None guidance_scale: float = 4.5 num_images_per_prompt: typing.Optional[int] = 1 height: typing.Optional[int] = None width: typing.Optional[int] = None eta: float = 0.0 generator: typing.Union[torch._C.Generator, typing.List[torch._C.Generator], NoneType] = None latents: typing.Optional[torch.Tensor] = None prompt_embeds: typing.Optional[torch.Tensor] = None prompt_attention_mask: typing.Optional[torch.Tensor] = None negative_prompt_embeds: typing.Optional[torch.Tensor] = None negative_prompt_attention_mask: typing.Optional[torch.Tensor] = None output_type: typing.Optional[str] = 'pil' return_dict: bool = True callback: typing.Optional[typing.Callable[[int, int, torch.Tensor], NoneType]] = None callback_steps: int = 1 clean_caption: bool = True use_resolution_binning: bool = True max_sequence_length: int = 120 **kwargs ) → ImagePipelineOutput 或 tuple
引數
- 提示 (
str
或List[str]
, 可選) — 用於引導影像生成的提示。如果未定義,則必須傳入prompt_embeds
。 - 負面提示 (
str
或List[str]
, 可選) — 不用於引導影像生成的提示。如果未定義,則必須傳入negative_prompt_embeds
。當不使用引導時(即,如果guidance_scale
小於1
時被忽略),則忽略。 - num_inference_steps (
int
, 可選, 預設為 100) — 去噪步數。更多去噪步數通常會帶來更高質量的影像,但推理速度會變慢。 - timesteps (
List[int]
, 可選) — 用於去噪過程的自定義時間步,適用於其set_timesteps
方法支援timesteps
引數的排程器。如果未定義,將使用傳遞num_inference_steps
時的預設行為。必須按降序排列。 - sigmas (
List[float]
, 可選) — 用於去噪過程的自定義 sigmas,適用於其set_timesteps
方法支援sigmas
引數的排程器。如果未定義,將使用傳遞num_inference_steps
時的預設行為。 - 引導比例 (
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,對其他排程器將被忽略。 - 生成器 (
torch.Generator
或List[torch.Generator]
, 可選) — 一個或多個 torch 生成器,用於使生成確定性。 - 潛在變數 (
torch.Tensor
, 可選) — 預生成的噪聲潛在變數,從高斯分佈中取樣,用作影像生成的輸入。可用於透過不同提示調整相同生成。如果未提供,將透過使用提供的隨機generator
取樣生成一個潛在變數張量。 - prompt_embeds (
torch.Tensor
, 可選) — 預生成的文字嵌入。可用於輕鬆調整文字輸入,例如 提示加權。如果未提供,將從prompt
輸入引數生成文字嵌入。 - prompt_attention_mask (
torch.Tensor
, 可選) — 預生成的文字嵌入的注意力掩碼。 - negative_prompt_embeds (
torch.Tensor
, 可選) — 預生成的負面文字嵌入。對於 PixArt-Alpha,此負面提示應為 ""。如果未提供,負面提示嵌入將從negative_prompt
輸入引數生成。 - negative_prompt_attention_mask (
torch.Tensor
, 可選) — 預生成的負面文字嵌入的注意力掩碼。 - output_type (
str
, 可選, 預設為"pil"
) — 生成影像的輸出格式。選擇 PIL:PIL.Image.Image
或np.array
。 - return_dict (
bool
, 可選, 預設為True
) — 是否返回~pipelines.stable_diffusion.IFPipelineOutput
而非普通元組。 - 回撥 (
Callable
, 可選) — 在推理過程中每callback_steps
步都會呼叫的函式。該函式將使用以下引數呼叫:callback(step: int, timestep: int, latents: torch.Tensor)
。 - callback_steps (
int
, 可選, 預設為 1) — 呼叫callback
函式的頻率。如果未指定,則在每個步驟都會呼叫回撥函式。 - clean_caption (
bool
, 可選, 預設為True
) — 是否在建立嵌入之前清理標題。需要安裝beautifulsoup4
和ftfy
。如果未安裝依賴項,則將從原始提示建立嵌入。 - use_resolution_binning (
bool
預設為True
) — 如果設定為True
,則首先使用ASPECT_RATIO_1024_BIN
將請求的高度和寬度對映到最接近的解析度。生成的潛在影像解碼為影像後,會將其大小調整回請求的解析度。對於生成非方形影像非常有用。 - max_sequence_length (
int
預設為 120) — 與prompt
一起使用的最大序列長度。
返回
ImagePipelineOutput 或 tuple
如果 return_dict
為 True
,則返回 ImagePipelineOutput,否則返回一個 tuple
,其中第一個元素是生成的影像列表
呼叫管道進行生成時呼叫的函式。
示例
>>> import torch
>>> from diffusers import PixArtAlphaPipeline
>>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-XL-2-512x512" too.
>>> pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-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: typing.Union[str, typing.List[str]] do_classifier_free_guidance: bool = True negative_prompt: str = '' num_images_per_prompt: int = 1 device: typing.Optional[torch.device] = None prompt_embeds: typing.Optional[torch.Tensor] = None negative_prompt_embeds: typing.Optional[torch.Tensor] = None prompt_attention_mask: typing.Optional[torch.Tensor] = None negative_prompt_attention_mask: typing.Optional[torch.Tensor] = None clean_caption: bool = False max_sequence_length: int = 120 **kwargs )
引數
- prompt (
str
或List[str]
, 可選) — 待編碼的提示詞 - negative_prompt (
str
或List[str]
, 可選) — 不用於引導影像生成的提示詞。如果未定義,則必須傳遞negative_prompt_embeds
。在使用無分類器引導時(即,如果guidance_scale
小於1
),此引數將被忽略。對於 PixArt-Alpha,此引數應為“”。 - do_classifier_free_guidance (
bool
, 可選, 預設為True
) — 是否使用無分類器引導 - 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
, 預設為 120) — 用於提示詞的最大序列長度。
將提示編碼為文字編碼器隱藏狀態。