[Core] Add PAG support for PixArtSigma #8921

sayakpaul · 2024-07-21T07:53:42Z

What does this PR do?

Part of #8785.

I think PixArt Sigma is way more popular than PixArt Alpha (longer sequence length, better quality in general, etc.), it makes sense to just add PAG to PixArt Sigma. If there's a need, I think we could open it up to community.

TODO

Tests
Docs

Code

from diffusers import PixArtSigmaPAGPipeline
import torch 

pipe = PixArtSigmaPAGPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", 
    torch_dtype=torch.float16, 
    pag_applied_layers=[14]
).to("cuda")

image = pipe(
    "A small cactus with a happy face in the Sahara desert.", 
    guidance_scale=0.0, 
    pag_scale=2.0, 
    generator=torch.manual_seed(0)
).images[0]

A bit of ablation study is in the comment below.

@asomoza would appreciate it if you could run some experiments with it :)

@sunovivid for awareness.

HuggingFaceDocBuilderDev · 2024-07-21T07:59:24Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

sayakpaul · 2024-07-21T08:02:52Z

Ablation

Code

import torch
import argparse
from diffusers import PixArtSigmaPAGPipeline, PixArtSigmaPipeline

def load_pipeline(args):
    if args.pag:
        pipe = PixArtSigmaPAGPipeline.from_pretrained(
            "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16, pag_applied_layers=args.pag_applied_layers
        ).to("cuda")
    else:
        pipe = PixArtSigmaPipeline.from_pretrained("PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16).to("cuda")
    return pipe 

def run_pipeline(pipe, args):
    if args.pag:
        image = pipe(
            args.prompt, guidance_scale=args.cfg, pag_scale=args.pag_scale, generator=torch.manual_seed(0)
        ).images[0]
    else:
        image = pipe(args.prompt, generator=torch.manual_seed(0)).images[0]
    
    img_name = "_".join(args.prompt.split(" "))
    if args.pag:
        img_name += "_pag"
        pag_applied_layers = list(map(str, args.pag_applied_layers))
        pag_applied_layers = "_".join(pag_applied_layers)
        img_name += f"_cfg@{args.cfg}_pg@{args.pag_scale}_layers@{pag_applied_layers}"
    img_name += ".png"
    image.save(img_name)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pag", type=int, default=0)
    parser.add_argument("--prompt", type=str, default="A small cactus with a happy face in the Sahara desert.")
    parser.add_argument("--cfg", type=float, default=0.0)
    parser.add_argument("--pag_scale", type=float, default=4.0)
    parser.add_argument("--pag_applied_layers", metavar="N", type=int, nargs="*", default=[14])
    args = parser.parse_args()

    pipe = load_pipeline(args)
    run_pipeline(pipe, args)

Results

Prompt: A small cactus with a happy face in the Sahara desert

W/O PAG	PAG

Prompt: Astronaut on Mars During sunset

W/O PAG	PAG_CFG_0_PG_2_LAYERS_14	PAG_CFG_0_PG_2_LAYERS_13_14	PAG_CFG_2_PG_2_LAYERS_14

sayakpaul · 2024-07-21T08:04:55Z

src/diffusers/pipelines/pag/pag_utils.py

@@ -258,3 +258,91 @@ def pag_attn_processors(self):
            if proc.__class__ in (PAGCFGIdentitySelfAttnProcessor2_0, PAGIdentitySelfAttnProcessor2_0):
                processors[name] = proc
        return processors
+
+
+class PixArtPAGMixin(PAGMixin):


Let me know if I should copy-paste the other methods from PAGMixin and not make it a subclass of PAGMixin.

PixArtPAGMixin differs only in terms of the methods I implemented here.

I'm also curious about this. subclass looks neater than copy-paste.

I would prefer not to have a Mixin nested into another Mixin if it's possible
we can either

copy-paste the common methods from PAGMixin (maybe we can rename PAGMixin to SDPAGMixin in that case)

we keep a common PAGMixin and move these methods that're potentially model-specific to pipeline methods

I will go with the first one. I will defer any refactoring to @yiyixuxu if she feels the need to do one.

Done in 27fceb3. LMK.

cc @a-r-r-o-w here too, I think we can reuse this one for hunyuan too

yiyixuxu · 2024-07-23T00:51:50Z

maybe we can try the robot prepare meal example? https://huggingface.co/docs/diffusers/main/en/using-diffusers/pag#general-tasks

I cannot tell if there is an improvement or not in the example provided (but I'm also not very good at looking at this, so cc @asomoza for help here too)

sayakpaul · 2024-07-23T02:12:33Z

Yeah I am not good at that either. So, requested @asomoza to chime in :D

Here are the insect cooking meal results.

No PAG	CFG_1_PG_3_L14	CFG_4.5_PG_3_L13_14_15	CFG_4.5_PG_3_L14

CFG_1_PG_3_L14 means we are applying PAG with the following config:

CFG of 1.0
PAG scale of 3.0
pag_applied_layers is set to [14]

asomoza · 2024-07-23T03:42:53Z

This one is the most interesting one of the new models, it definitely helps with the generation, it cleans the image and fixes the errors, but as with kolors, it takes away some details, so it depends on what you want.

For example in this image, it fixes the hands of the robot and the shape of the bowls:

W/O PAG	PAG

Also I found the time to play with the model and it's a really good model, I'm impressed with the generations it can do.

Anyway, what I found interesting is how the PAG layers behave, they affect a lot of the generation, for example:

L1	L2	L6

L18	L19	L21

As with SDXL, I can see some relation of the layers with the generations but I'll need to play more with it to be sure. For example the 19 and 21 layers seems to affect more the background and some others seems to affect more the shape or the colors. Also the 1 and 6 layers can be used for the composition with a second pass over them.

sayakpaul · 2024-07-23T04:37:38Z

Thank you very much, Alvaro! So, what I am hearing is that it makes sense to have PAG supported for PixArt Sigma, yeah?

Cc @lawrence-cj too you might find it nice.

src/diffusers/pipelines/pag/pag_utils.py

yiyixuxu · 2024-07-24T01:45:21Z

@asomoza very nice! thanks!

sayakpaul · 2024-07-26T01:31:30Z

@yiyixuxu can I add test and docs and get it ready for final review?

yiyixuxu · 2024-07-27T07:35:02Z

@sayakpaul sure ! looks good to me
let's merge this soon

sayakpaul · 2024-07-27T12:51:15Z

tests/pipelines/pag/test_pag_pixart_sigma.py

+    params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"})
+    params = set(params)
+    params.remove("cross_attention_kwargs")


LMK if this needs to be handled differently. PixArt doesn't have cross_attention_kwargs and doesn't need to have yet.

sayakpaul · 2024-07-27T12:51:48Z

tests/pipelines/pag/test_pag_pixart_sigma.py

+    # Because the PAG PixArt Sigma has `pag_applied_layers`.
+    # Also, we shouldn't be doing `set_default_attn_processor()` after loading
+    # the pipeline with `pag_applied_layers`.


Let me know if this and subsequent tests should be handled differently.

looks good!

sayakpaul · 2024-07-27T17:25:58Z

@yiyixuxu I have left some questions for you regarding the tests. LMK.

yiyixuxu · 2024-08-01T23:36:30Z

@sayakpaul looks good to me! feel free to merge once the conflicts are resolved

* feat: add pixart sigma pag. * inits. * fixes * fix * remove print. * copy paste methods to the pixart pag mixin * fix-copies * add documentation. * add tests. * remove correction file. * remove pag_applied_layers * empty

sayakpaul added 5 commits July 21, 2024 09:13

feat: add pixart sigma pag.

e73b269

inits.

67e6442

fixes

9dda0b5

fix

124cf15

remove print.

6123399

sayakpaul requested a review from yiyixuxu July 21, 2024 07:53

sayakpaul commented Jul 21, 2024

View reviewed changes

Merge branch 'main' into pag-pixart-sigma

502cdb0

sayakpaul added 3 commits July 23, 2024 07:46

copy paste methods to the pixart pag mixin

27fceb3

fix-copies

0922fc6

Merge branch 'main' into pag-pixart-sigma

af5218c

Merge branch 'main' into pag-pixart-sigma

a39ab9a

yiyixuxu reviewed Jul 24, 2024

View reviewed changes

src/diffusers/pipelines/pag/pag_utils.py Show resolved Hide resolved

sayakpaul added 2 commits July 24, 2024 08:37

Merge branch 'main' into pag-pixart-sigma

a996152

Merge branch 'main' into pag-pixart-sigma

df8ff85

Merge branch 'main' into pag-pixart-sigma

e16f356

sayakpaul added 5 commits July 27, 2024 17:04

Merge branch 'main' into pag-pixart-sigma

7e96451

add documentation.

c239549

add tests.

3540e06

remove correction file.

91c7603

remove pag_applied_layers

e63f7cd

sayakpaul commented Jul 27, 2024

View reviewed changes

empty

ddfd451

sayakpaul added 2 commits July 30, 2024 08:58

Merge branch 'main' into pag-pixart-sigma

dd02c4f

Merge branch 'main' into pag-pixart-sigma

f044e7e

resolve conflicts.

e88a2d6

sayakpaul merged commit 7b98c4c into main Aug 2, 2024
18 checks passed

sayakpaul deleted the pag-pixart-sigma branch August 2, 2024 01:57

yiyixuxu added the PAG label Sep 4, 2024

[Core] Add PAG support for PixArtSigma #8921

[Core] Add PAG support for PixArtSigma #8921

Uh oh!

Conversation

sayakpaul commented Jul 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

TODO

Code

Uh oh!

HuggingFaceDocBuilderDev commented Jul 21, 2024

Uh oh!

sayakpaul commented Jul 21, 2024

Ablation

Code

Results

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yiyixuxu Jul 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yiyixuxu commented Jul 23, 2024

Uh oh!

sayakpaul commented Jul 23, 2024

Uh oh!

asomoza commented Jul 23, 2024

Uh oh!

sayakpaul commented Jul 23, 2024

Uh oh!

Uh oh!

yiyixuxu commented Jul 24, 2024

Uh oh!

sayakpaul commented Jul 26, 2024

Uh oh!

yiyixuxu commented Jul 27, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul commented Jul 27, 2024

Uh oh!

yiyixuxu commented Aug 1, 2024

Uh oh!

Uh oh!

Uh oh!

sayakpaul commented Jul 21, 2024 •

edited

Loading

yiyixuxu Jul 23, 2024 •

edited

Loading