[core] Refactor condition embedders (and others) modules from pipelines#13781
[core] Refactor condition embedders (and others) modules from pipelines#13781sayakpaul wants to merge 27 commits into
pipelines#13781Conversation
|
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. |
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| import torch |
There was a problem hiding this comment.
Open to suggestions on where to place these other than others.
| @@ -0,0 +1,476 @@ | |||
| # Copyright 2025 The ACE-Step Team and The HuggingFace Team. All rights reserved. | |||
There was a problem hiding this comment.
Going by the implementation of the module, it encodes into a tokenized representation (not the existing sense of tokenization done in language models). So, decided keep them under autoencoders.
| @@ -0,0 +1,192 @@ | |||
| # Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. | |||
There was a problem hiding this comment.
Open to moving them elsewhere. Since upsampling could be interpreted as decoding to a higher resolution, I thought of keeping them under autoencoders.
|
|
||
| # Reroutes pretrained loads past pipeline-local deprecation shims onto the canonical | ||
| # top-level diffusers export. | ||
| _RELOCATED_PIPELINE_CLASSES: dict[tuple[str, str], tuple[str, str]] = { |
There was a problem hiding this comment.
This is so that we don't hit deprecation warnings.
| @@ -0,0 +1,93 @@ | |||
| import sys | |||
There was a problem hiding this comment.
🤗 Serge says:
Overall
This is a well-structured refactoring that moves model-level components out of pipelines/ into models/. The deprecation shim pattern (subclass with __new__ warning) is clean and preserves backward compatibility. However, there are a few issues worth addressing.
Correctness
-
Incomplete
_RELOCATED_PIPELINE_CLASSESmapping. The remap dict inpipeline_loading_utils.pyonly covers 6 of the relocated classes. The remaining ones (LTX2VocoderWithBWE,AceStepConditionEncoder,AceStepAudioTokenizer,AceStepAudioTokenDetokenizer,CLIPImageProjection,StableUnCLIPImageNormalizer,ReduxImageEncoder,ShapERenderer, and the LTX latent upsamplers) are missing. When a user loads a pipeline whosemodel_index.jsonreferences these classes via the old pipeline module name (e.g.["ace_step", "AceStepConditionEncoder"]), the loader will resolve them through the deprecation shim and emit spurious deprecation warnings during normal pipeline usage. The remap dict should include all relocated classes so thatfrom_pretrainedon a pipeline bypasses the shim entirely. (This is visible in the full PR but the dict definition is not in this chunk — noting for completeness.) -
Pre-existing bug in
vocoder_ltx2.pyline 263:act_fn = nn.LeakyReLU(...)should beact = nn.LeakyReLU(...). This reassigns the string parameteract_fnto annn.Moduleand leavesactundefined in theelsebranch. This is a bug carried over from the original file, not introduced by this PR, but since the code is being relocated this is a good opportunity to fix it.
Style / Minor
vocoder_ltx2.pyandtext_connector_ltx2.pyare missing license headers, while all other new files in this PR have them. Likely carried over from the originals.
Ephemeral content
scripts/smoke_test_relocated_pretrained.py— the PR description says this will be deleted after merge. That's fine, but make sure it doesn't land onmain.
Dead code analysis (advisory)
The ResBlock class in vocoder_ltx2.py (line 214) shadows the ResBlock in latent_upsampler_ltx.py / latent_upsampler_ltx2.py within the same autoencoders package. They are different classes (1D conv vs 2D/3D conv). This is not a bug since they live in separate modules, but the name collision could confuse readers browsing the autoencoders package. The latent_upsampler_ltx2.py version is already annotated with # Copied from which helps.
Note: review finished early after hitting the input-token budget; 3 remaining diff chunk(s) were not reviewed.
model: claude-opus-4-6 · 25 LLM turns · 31 tool calls · 148.2s · 2093370 in / 6527 out tokens
What does this PR do?
We have a bunch of model-level components under
src/diffusers/pipelinesmodule. This PR factors them out of thepipelinesmodule and places them undersrc/diffusers/modelsmodule strategically.The main changes are in
src/pipelines/pipeline_loading_utils.pywhich basically remaps to the components that were moved tosrc/diffusersmodels/. Open to other ways of handling it.For reviewers
Main changes can be summarized as follows:
audio_tokenizer_ace_step.py) tosrc/diffusers/modelsand deprecate gracefully.autoencodersandcondition_embedders. When I wasn't sure, I simply placed them underothers. No strong opinions here.scripts/smoke_test_relocated_pretrained.pyis for transient testing until the PR is merged. Will delete it after. I have run the tests and models and pipelines load as expected.tests/models/test_relocated_class_loading.pywhich will check if deprecation warnings are raised and will also help to remove those deprecation paths once the cycle is over.tests/pipelines/test_relocated_class_loading.py. It ensures that no deprecation warnings for the refactored modules are raised when loading the affected pipelines.