• Hameer Abbasi's avatar
    [LoRA] Add LoRA support to AuraFlow (#10216) · 9352a5ca
    Hameer Abbasi authored
    
    
    * Add AuraFlowLoraLoaderMixin
    
    * Add comments, remove qkv fusion
    
    * Add Tests
    
    * Add AuraFlowLoraLoaderMixin to documentation
    
    * Add Suggested changes
    
    * Change attention_kwargs->joint_attention_kwargs
    
    * Rebasing derp.
    
    * fix
    
    * fix
    
    * Quality fixes.
    
    * make style
    
    * `make fix-copies`
    
    * `ruff check --fix`
    
    * Attept 1 to fix tests.
    
    * Attept 2 to fix tests.
    
    * Attept 3 to fix tests.
    
    * Address review comments.
    
    * Rebasing derp.
    
    * Get more tests passing by copying from Flux. Address review comments.
    
    * `joint_attention_kwargs`->`attention_kwargs`
    
    * Add `lora_scale` property for te LoRAs.
    
    * Make test better.
    
    * Remove useless property.
    
    * Skip TE-only tests for AuraFlow.
    
    * Support LoRA for non-CLIP TEs.
    
    * Restore LoRA tests.
    
    * Undo adding LoRA support for non-CLIP TEs.
    
    * Undo support for TE in AuraFlow LoRA.
    
    * `make fix-copies`
    
    * Sync with upstream changes.
    
    * Remove unneeded stuff.
    
    * Mirror `Lumina2`.
    
    * Skip for MPS.
    
    * Address review comments.
    
    * Remove duplicated code.
    
    * Remove unnecessary code.
    
    * Remove repeated docs.
    
    * Propagate attention.
    
    * Fix TE target modules.
    
    * MPS fix for LoRA tests.
    
    * Unrelated TE LoRA tests fix.
    
    * Fix AuraFlow LoRA tests by applying to the right denoiser layers.
    Co-authored-by: default avatarAstraliteHeart <81396681+AstraliteHeart@users.noreply.github.com>
    
    * Apply style fixes
    
    * empty commit
    
    * Fix the repo consistency issues.
    
    * Remove unrelated changes.
    
    * Style.
    
    * Fix `test_lora_fuse_nan`.
    
    * fix quality issues.
    
    * `pytest.xfail` -> `ValueError`.
    
    * Add back `skip_mps`.
    
    * Apply style fixes
    
    * `make fix-copies`
    
    ---------
    Co-authored-by: default avatarWarlord-K <warlordk28@gmail.com>
    Co-authored-by: default avatarhlky <hlky@hlky.ac>
    Co-authored-by: default avatarSayak Paul <spsayakpaul@gmail.com>
    Co-authored-by: default avatarAstraliteHeart <81396681+AstraliteHeart@users.noreply.github.com>
    Co-authored-by: default avatargithub-actions[bot] <github-actions[bot]@users.noreply.github.com>
    9352a5ca
lora_pipeline.py 273 KB