PixArt Sigma PEFT LoRA loader support

[bghira]

Is your feature request related to a problem? Please describe.

Currently, the upstream PixArt trainer does this:

 transformer = get_peft_model(transformer, lora_config)
 if args.mixed_precision == "fp16":
 # only upcast trainable parameters (LoRA) into fp32
 cast_training_params(transformer, dtype=torch.float32)

 transformer.print_trainable_parameters()

 # 10. Handle saving and loading of checkpoints
 # `accelerate` 0.16.0 will have better support for customized saving
 if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
 # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
 def save_model_hook(models, weights, output_dir):
 if accelerator.is_main_process:
 transformer_ = accelerator.unwrap_model(transformer)
 lora_state_dict = get_peft_model_state_dict(transformer_, adapter_name="default")
 StableDiffusionPipeline.save_lora_weights(os.path.join(output_dir, "transformer_lora"), lora_state_dict)
 # save weights in peft format to be able to load them back
 transformer_.save_pretrained(output_dir)

 for _, model in enumerate(models):
 # make sure to pop weight so that corresponding model is not saved again
 weights.pop()

Describe the solution you'd like.

I would like for the PixArtSigmaPipeline to have the necessary mix-ins and inference support instead.

Describe alternatives you've considered.

I have considered avoiding the pipeline methods and using workarounds like upstream does, but instead I would prefer consistency for user experience.

[bghira]

#9887 as well

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[bghira]

not stale

[M0315G]

Hey @yiyixuxu , Is this issue still open for contributions? I'd like to give this a try.

[yiyixuxu]

@M0315G yes! please

[M0315G]

Perfect! I'll go through the contribution guidelines, anything else you'd want me to review before I start?

[M0315G]

Hey @yiyixuxu , went through the files and this is what I understood (correct me if I'm wrong):

1. Add upscaling and downscaling support for LoRA in PixArtTransformer2DModel model file.
2. Create a PixArtLoraLoaderMixin class with methods to load and unload the Loras. Specifically
   --> load_lora_into_transformer
   --> load_lora_into_text_encoder

I'm still going through rest of them like fuse_lora and unfuse_lora.

[M0315G]

Soft ping @yiyixuxu !

[yiyixuxu]

you can maybe follow this PR as an exanple to add lora support #10234
but I think we should start with an example lora checkpoint we want to support first

maybe @bghira can provide an example?

cc @sayakpaul here, do you have additional guidance you want to provide?

[sayakpaul]

Yes, having an example LoRA with supported added would be a good first step. Besides that, I don't have nothing much to add.

[M0315G]

Ping @bghira

[bghira]

@M0315G
i don't think anything special is required for Diffusers+PEFT style LoRAs like SimpleTuner creates. these are nearly universally loadable by Diffusers. however, probably a pixart sigma LoRA from CivitAI would be helpful to target third party key layouts like Kohya or Onetrainer. i'm not sure which key layouts exist, it's possible to use the CivitAI API and mass-download Pixart LoRA (and Lycoris) and collate any common key layout elements, so that all of them are supported by the LoRA state key dict conversion methods.