Skip to content

Fix shared LoRA component weight broadcasting - #15128

Open
gorkadabo1 wants to merge 1 commit into
huggingface:mainfrom
gorkadabo1:fix-lora-component-weight-broadcast
Open

gorkadabo1 wants to merge 1 commit into
huggingface:mainfrom
gorkadabo1:fix-lora-component-weight-broadcast

Conversation

@gorkadabo1

Copy link
Copy Markdown

What does this PR do?

Fixes #15127. Approved by @sayakpaul.

When set_adapters() receives a single component-weight dictionary for multiple LoRA adapters, only the first adapter gets the requested weight. The others use the default weight of 1.0 without warning. This PR replaces pop() with get(), so the same dictionary can be used for the remaining adapters.

It also adds regression tests for broadcast and shared-reference dictionaries, zero weights, reversed adapter order, and nested UNet block weights. The tests also check that inference outputs match.

Validation

Python 3.12.1, PyTorch 2.14.1 (CPU), PEFT 0.21.2.

  • New regression tests: all 7 failed without the fix and passed with it.
  • Complete SDXL and Flux LoRA test classes: 77 passed, 110 warnings.
  • make style: passed.
  • make fix-copies: passed.
  • make quality: passed.
  • git diff --check: passed.

Tests used CPU float32 and small test models. I did not run GPU, mixed-precision, quantized, or full-library tests.

python -m pytest tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py::TestStableDiffusionXLPipelineLoRA tests/pipelines/flux/test_pipeline_flux.py::TestFluxPipelineLoRA -k 'shared_component_weights or shared_block_weights' -q --tb=short
python -m pytest tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py::TestStableDiffusionXLPipelineLoRA tests/pipelines/flux/test_pipeline_flux.py::TestFluxPipelineLoRA -q --tb=short

AI assistance and self-review

I used Codex to help with the implementation, testing, and code review. The final self-review report, including findings, validation results, and limitations, is included below.

Final codex self-review report
DIFFUSERS #15127 — FINAL SELF-REVIEW (SIMPLIFIED TESTS)
  Base: aaf57d03a3edfbd9de457d2872417a92f149a91d
  Branch: fix-lora-component-weight-broadcast
  Patch SHA-256: a940c6b9f265fcffb40afac6e4ba286f892752fc96bdb63db1275340a11c5339

  Applied .ai/skills/self-review/SKILL.md to the complete uncommitted git diff HEAD,
  using .ai/references/review-rules.md and the applicable code-style, shared-test,
  pipeline/model, nested-scale, and numerical guidance. This is a Codex self-review.

  Blocking findings: none identified.

  Non-blocking suggestion:
  1. Existing docstring underspecifies supported weight forms.
     src/diffusers/loaders/lora_base.py:709 still describes list weights and UNet
     only. Per review-rules.md (Documentation impact), optionally clarify component
     dictionary broadcasting. Left unchanged to keep this bug fix focused; no new
     API is introduced. No agent instructions or configuration were modified.

  Correctness:
  - Production fix is unchanged: weights.pop(component, None) becomes
    weights.get(component, None). No later code relies on consumed component keys.
  - Scalar/list handling, validation, filtering, optional components, ownership
    warnings, and None defaults keep their existing control flow.
  - Both ModelMixin.set_adapters and text-encoder helper paths were checked.
    Nested UNet expansion deep-copies its input; activation reads weights.
    The existing outer deepcopy remains in place.

  Test proportionality and conventions:
  - Two shared-mixin test methods, with three general parameter pairs: broadcast
    0.5, broadcast 0.0, and repeated-reference list 0.5. Executed for SDXL and Flux.
  - One broadcast nested UNet block-weight case in SDXL. Seven retained cases
    across the two validated classes; each fails against original upstream code.
  - Independent-dictionary reference outputs, original caller configuration,
    reversed adapter order, and actual inference comparisons are retained.
  - Removed unit-weight cases, redundant alias/zero and nested alias combinations,
    and repeated scalar-reference inference. No bespoke fixtures or test hooks.
  - Existing real tiny pipeline builders, nonzero adapter helper, seeded run_pipe,
    and tensor assertions follow testing.md. Shared tests also collect for other
    compatible pipelines; only SDXL and Flux classes were executed here.

  Validation:
  - Baseline production file verified byte-for-byte against HEAD in a separate
    worktree; baseline imported its own src/diffusers package. Same simplified tests.
  - Baseline regressions: 7 failed, 70 deselected, 19 warnings in 6.55s (exit 1).
  - Fixed regressions: 7 passed, 70 deselected, 22 warnings in 4.80s (exit 0).
  - Complete SDXL and Flux LoRA classes: 77 passed, 110 warnings in 40.31s (exit 0).
    This comprises 70 pre-existing cases and seven retained new cases.
  - make style, make fix-copies, make quality, git diff --check: all exit 0.
  - make style collapsed the parameter decorator without changing its AST. Final
    tests have identical AST to the baseline/fixed runs. The production file stayed
    byte-identical. An unrelated generated loader import reorder was restored.

  Dead-code advisory:
  Path:line | Status | Reason
  src/diffusers/loaders/lora_base.py:780 | Used | Pipeline component-weight lookup.
  tests/pipelines/testing_utils/lora.py:705 | Used | Six collected cases across SDXL/Flux.
  tests/pipelines/testing_utils/lora.py:1227 | Used | One collected nested UNet case in SDXL.
  No dead code identified; no new models, dependencies, or fallback branches.

  Verdict: READY for human review.
  Fix before submitting: no blocking findings or dead code identified.
  Leave for maintainer review: optional existing docstring clarification above.
  Limits: CPU float32 tiny fixtures; no GPU, mixed precision, quantization,
  full-size model, offloading class, or entire-library validation.
  Final self-review notes belong in the eventual PR description/comment, outside
  its code diff. The PR draft includes the findings and limitations above.

Before submitting

Who can review?

@BenjaminBossan, as requested by @sayakpaul.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

set_adapters applies component weights only to the first adapter when broadcasting a dictionary

1 participant