Diffusers代码学习: PixArt-Σ

PixArt-Σ，一种扩散变压器模型（DiT），能够直接生成4K分辨率的图像。

1. 使用 Transformer 主干网（而不是 UNet）进行去噪。因此，具有与 DiT 类似的架构。

2. 使用从 T5 计算的文本条件进行训练的。这方面使管道能够更好地遵循具有复杂细节的复杂文本提示。

3. 擅长以不同的纵横比生成高分辨率图像。

4. 质量可与最先进的文生图系统相媲美。

5. 显示了生成超高分辨率图像的能力，例如 2048px 甚至 4K。

通过多项改进（VAE、数据集等），文生图模型可以从弱模型发展到更强的模型。

import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

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

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


import gc

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

del text_encoder
del pipe
flush()

pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-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()

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

Diffusers代码学习: PixArt-Σ

悦读