[WIP] PixArt-Sigma training pipeline

[kabachuha]

Current state

• Backbone (w/o external deps)
• Model save-loading (orig format)
• T5 support, T5's attention_mask carryover, load in 4bit
• T5 Text embeddings, attention mask caching on disc
• Wrap up Pixart as a NetworkTrainer for use in any train loops
• Add pixart blocks for lora/etc listings
• Combine T5 and SDXL vae in checkpoint when saving, recommended by FurkanGozukara
• Model inference for sampling
• Do test launch on base and lora/etc to test compat and debug
• Test aspect-ratio conditioning
• Diffusers format save mode and other leftover TODOs
• Setup a ShareCaption/CogVLM2/LLaVA*/etc. InternLM-XComposer2-4KHD-based multimodal prompt enhancer
• ControlNet-Transformer support for training and inference (waiting for PA-Sigma's repo code release)

I'm very excited to work with PixArt and its great size/prompt adherence ratio, in addition to the awesome lora techniques in this repo, so if it goes well it should start working in a couple of days

Addresses #979

PixArt repo: https://github.com/PixArt-alpha/PixArt-sigma

I'm very likely going to edit this post with updates, comments, pictures quite often

[FurkanGozukara]

awesome will this have single safetensors file ? PixArt could be the future if SD3 never get released

[kabachuha]

awesome will this have single safetensors file

Yes, ofc the diffusion transformer is saved in safetensors format, if the safetensors option is specified.

As to embedding T5 and the SD vae inside along with it, it may be an option, but not sure about it's being practical, as the original PA format loads them from the separated hf files and the existing workflows such as in comfy load them from the different subfolders (T5 and SDXL VAE) as well

PixArt could be the future if SD3 never get released

Yes the SD3's future is clouded. Additionally, it's parameters count is much bigger than PA-Sigma with comparable quality and prompt adherence as we see from the api examples, meaning it's more accessible and faster, enabling easier community fine-tuning. And regardless of the things, we have PA now, and we can experiment with different things and techniques on Diffusion Transformers, that can be handy when another transformer-like diffusion is released

[FurkanGozukara]

@kabachuha can we make it include text encoder and vae as well like single sdxl checkpoints and load from it?

I plan to make a standalone gradio app and maybe auto1111 adds support later

[kabachuha]

Very well, I see your point :)

[raulc0399]

hi @kabachuha just a fyi:

the sigma and alpha models have scripts for HF Lora training.
https://github.com/PixArt-alpha/PixArt-sigma/blob/master/train_scripts/train_pixart_lora_hf.py
https://github.com/PixArt-alpha/PixArt-alpha/blob/master/train_scripts/train_pixart_lora_hf.py

there are also PRs for text-encoder training in parallel.
if you want to integrate those i can provide help with that.

[kabachuha]

Yes, I know. But it's a kohya repo, and the lora modules and the training scrips have different styles, so need to suit them :)

If you'd like to help, it would be nice (I can add you to the fork, so you may be able to take over the pr, as I have load this week. Dm or comment here)

[raulc0399]

@kabachuha yes, please add me i will start next week.
do you also have an overview of what has been done so far and what needs to be done?

[kabachuha]

Added :) And see the list

[raulc0399]

@kabachuha so the still open todos are following. looks like mostly testing?

• Combine T5 and SDXL vae in checkpoint when saving, recommended by FurkanGozukara
• Do test launch on base and lora/etc to test compat and debug
• Test aspect-ratio conditioning
• Diffusers format save mode and other leftover TODOs
• Setup a ShareCaption/CogVLM2/LLaVA*/etc. InternLM-XComposer2-4KHD-based multimodal prompt enhancer

do you have any more details on the "Combine T5 and SDXL vae" and "other leftover TODOs" ?

[AtomisteBX]

Any news? SD3 is a disappointment so maybe there should be more focus towards training Pixart Sigma...

[DanPli]

"Combine T5 and SDXL vae in checkpoint when saving" Is not a good idea, if the T5/vae (mostly the T5 checkpoint) gets multiplied by this, people will just go use other training scripts that make alot smaller checkpoints, and rightfully so. Storage space reqs would explode with increasing amounts of custom networks trained or downloaded. You should have one T5 and a VAE separate for not only this reason, also to be able to modularly approach and substitute other future Text or Autoencoders and combine them freely with any finetune/network you have.