From 3ab185fea4cb7e4f023f0b0cb27902df8781eeb9 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 23 Oct 2026 05:31:57 +0000 Subject: [PATCH 1/5] move latent upscalers and conditionsn embedders out of pipelines --- docs/source/en/api/pipelines/ltx2.md | 10 +- docs/source/en/api/pipelines/ltx_video.md | 2 +- scripts/convert_ace_step_to_diffusers.py | 2 +- scripts/convert_ltx2_to_diffusers.py | 4 +- scripts/convert_ltx_to_diffusers.py | 2 +- scripts/smoke_test_relocated_pretrained.py | 109 ++++ src/diffusers/__init__.py | 31 +- src/diffusers/loaders/single_file.py | 6 + src/diffusers/models/__init__.py | 23 +- .../models/condition_embedders/__init__.py | 7 + .../condition_embedder_ace_step.py | 477 ++++++++++++++++ .../condition_embedder_audioldm2.py | 142 +++++ .../condition_embedder_clip_image.py | 29 + .../condition_embedder_flux_redux.py | 46 ++ .../condition_embedder_ltx2.py | 478 ++++++++++++++++ .../condition_embedder_stable_audio.py | 156 ++++++ .../condition_embedder_stable_audio_3.py | 109 ++++ .../models/latent_upscaler/__init__.py | 2 + .../latent_upscaler/latent_upscaler_ltx.py | 186 +++++++ .../latent_upscaler/latent_upscaler_ltx2.py | 285 ++++++++++ .../modular_pipelines/echo/encoders.py | 2 +- .../modular_pipelines/ltx2/encoders.py | 17 +- .../pipelines/ace_step/modeling_ace_step.py | 477 ++-------------- .../pipelines/ace_step/pipeline_ace_step.py | 3 +- .../pipelines/audioldm2/modeling_audioldm2.py | 143 +---- .../pipelines/audioldm2/pipeline_audioldm2.py | 3 +- ...line_stable_diffusion_gligen_text_image.py | 2 +- src/diffusers/pipelines/flux/modeling_flux.py | 50 +- .../flux/pipeline_flux_prior_redux.py | 2 +- .../ltx/modeling_latent_upsampler.py | 189 +------ .../ltx/pipeline_ltx_latent_upsample.py | 2 +- src/diffusers/pipelines/ltx2/connectors.py | 512 ++---------------- .../pipelines/ltx2/latent_upsampler.py | 291 +--------- src/diffusers/pipelines/ltx2/pipeline_ltx2.py | 2 +- .../pipelines/ltx2/pipeline_ltx2_condition.py | 2 +- .../pipelines/ltx2/pipeline_ltx2_dfr.py | 2 +- .../ltx2/pipeline_ltx2_dfr_temporal_refine.py | 6 +- .../pipelines/ltx2/pipeline_ltx2_hdr_lora.py | 2 +- .../pipelines/ltx2/pipeline_ltx2_ic_lora.py | 2 +- .../ltx2/pipeline_ltx2_image2video.py | 2 +- .../ltx2/pipeline_ltx2_latent_upsample.py | 4 +- .../pipelines/pipeline_loading_utils.py | 27 + src/diffusers/pipelines/sana_wm/refiner.py | 6 +- .../stable_audio/modeling_stable_audio.py | 162 +----- .../stable_audio/pipeline_stable_audio.py | 2 +- .../stable_audio_3/modeling_stable_audio_3.py | 112 +--- .../stable_audio_3/pipeline_stable_audio_3.py | 2 +- .../pipeline_stable_audio_3_audio2audio.py | 2 +- .../pipeline_stable_audio_3_inpaint.py | 2 +- .../clip_image_project_model.py | 26 +- src/diffusers/utils/dummy_pt_objects.py | 150 ++++- .../dummy_torch_and_transformers_objects.py | 90 --- tests/models/test_relocated_class_loading.py | 85 +++ tests/pipelines/ace_step/test_ace_step.py | 3 +- .../pipelines/ltx/test_ltx_latent_upsample.py | 2 +- tests/pipelines/ltx2/test_ltx2_connectors.py | 2 +- tests/pipelines/ltx2/testing_utils.py | 4 +- .../stable_audio_3/test_stable_audio_3.py | 2 +- 58 files changed, 2570 insertions(+), 1930 deletions(-) create mode 100644 scripts/smoke_test_relocated_pretrained.py create mode 100644 src/diffusers/models/condition_embedders/condition_embedder_ace_step.py create mode 100644 src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py create mode 100644 src/diffusers/models/condition_embedders/condition_embedder_clip_image.py create mode 100644 src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py create mode 100644 src/diffusers/models/condition_embedders/condition_embedder_ltx2.py create mode 100644 src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py create mode 100644 src/diffusers/models/condition_embedders/condition_embedder_stable_audio_3.py create mode 100644 src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py create mode 100644 src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py create mode 100644 tests/models/test_relocated_class_loading.py diff --git a/docs/source/en/api/pipelines/ltx2.md b/docs/source/en/api/pipelines/ltx2.md index 2fa85c9e4af9..7bb173426340 100644 --- a/docs/source/en/api/pipelines/ltx2.md +++ b/docs/source/en/api/pipelines/ltx2.md @@ -39,7 +39,7 @@ Sample usage of text-to-video two stages pipeline import torch from diffusers import FlowMatchEulerDiscreteScheduler from diffusers.pipelines.ltx2 import LTX2Pipeline, LTX2LatentUpsamplePipeline -from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.models.latent_upscaler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES from diffusers.utils import encode_video @@ -125,7 +125,7 @@ Fastest two-stages generation pipeline using a distilled checkpoint. ```py import torch from diffusers.pipelines.ltx2 import LTX2Pipeline, LTX2LatentUpsamplePipeline -from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.models.latent_upscaler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES from diffusers.utils import encode_video @@ -203,7 +203,7 @@ You can use `LTX2ConditionPipeline` to specify image and/or video conditions at ```py import torch from diffusers import LTX2ConditionPipeline, LTX2LatentUpsamplePipeline -from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.models.latent_upscaler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES from diffusers.utils import encode_video @@ -569,7 +569,7 @@ Stage 1 runs at half the target resolution, the upsampler doubles it, and stage ```py import torch from diffusers.pipelines.ltx2 import LTX2Pipeline, LTX2LatentUpsamplePipeline -from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.models.latent_upscaler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES from diffusers.utils import encode_video @@ -1119,7 +1119,7 @@ from diffusers import ( LTX2LatentUpsamplePipeline, LTXEulerAncestralRFScheduler, ) -from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel +from diffusers.models.latent_upscaler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.utils import trim_canvas from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES diff --git a/docs/source/en/api/pipelines/ltx_video.md b/docs/source/en/api/pipelines/ltx_video.md index b7726d8f0c0b..c059094e772c 100644 --- a/docs/source/en/api/pipelines/ltx_video.md +++ b/docs/source/en/api/pipelines/ltx_video.md @@ -334,7 +334,7 @@ export_to_video(video, "output.mp4", fps=24) import torch from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition - from diffusers.pipelines.ltx.modeling_latent_upsampler import LTXLatentUpsamplerModel + from diffusers.models.latent_upscaler import LTXLatentUpsamplerModel from diffusers.utils import export_to_video, load_video pipeline = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.8-13B-distilled", dtype=torch.bfloat16) diff --git a/scripts/convert_ace_step_to_diffusers.py b/scripts/convert_ace_step_to_diffusers.py index 252f5517f91b..97cd76b6dfb1 100644 --- a/scripts/convert_ace_step_to_diffusers.py +++ b/scripts/convert_ace_step_to_diffusers.py @@ -330,6 +330,7 @@ def _rename_attn_keys(key: str) -> str: from transformers import AutoModel, AutoTokenizer from diffusers import ( + AceStepConditionEncoder, AceStepPipeline, AceStepTransformer1DModel, AutoencoderOobleck, @@ -338,7 +339,6 @@ def _rename_attn_keys(key: str) -> str: from diffusers.pipelines.ace_step import ( AceStepAudioTokenDetokenizer, AceStepAudioTokenizer, - AceStepConditionEncoder, ) # Drop metadata keys — they're re-populated by `save_pretrained` at save time. diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py index 33b91790ef1c..9fc3b81efe63 100644 --- a/scripts/convert_ltx2_to_diffusers.py +++ b/scripts/convert_ltx2_to_diffusers.py @@ -14,14 +14,14 @@ AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler, LTX2LatentUpsamplePipeline, + LTX2LatentUpsamplerModel, LTX2Pipeline, + LTX2TextConnectors, LTX2VideoDiffusionDecoderModel, LTX2VideoTransformer3DModel, ) from diffusers.pipelines.ltx2 import ( LTX2DurationHead, - LTX2LatentUpsamplerModel, - LTX2TextConnectors, LTX2Vocoder, LTX2VocoderWithBWE, ) diff --git a/scripts/convert_ltx_to_diffusers.py b/scripts/convert_ltx_to_diffusers.py index b4ae57f9c8f0..a59dab5916d0 100644 --- a/scripts/convert_ltx_to_diffusers.py +++ b/scripts/convert_ltx_to_diffusers.py @@ -15,7 +15,7 @@ LTXPipeline, LTXVideoTransformer3DModel, ) -from diffusers.pipelines.ltx.modeling_latent_upsampler import LTXLatentUpsamplerModel +from diffusers.models.latent_upscaler import LTXLatentUpsamplerModel def remove_keys_(key: str, state_dict: Dict[str, Any]): diff --git a/scripts/smoke_test_relocated_pretrained.py b/scripts/smoke_test_relocated_pretrained.py new file mode 100644 index 000000000000..c1589b53a664 --- /dev/null +++ b/scripts/smoke_test_relocated_pretrained.py @@ -0,0 +1,109 @@ +import os +import sys +import traceback +import warnings + +import torch + +import diffusers + + +# Stable Audio 3 is not published in diffusers format; convert it locally with +# `scripts/convert_stable_audio_3_to_diffusers.py` and point `SA3_CHECKPOINT` at the output directory. +SA3_CHECKPOINT = os.environ.get("SA3_CHECKPOINT", "stabilityai/stable-audio-3-medium") + + +# (class_name, repo, subfolder, kwargs) — subfolder=None loads from the repo root. +MODEL_CHECKPOINTS = [ + ("CLIPImageProjection", "anhnct/Gligen_Text_Image", "image_project", {}), + ("AudioLDM2ProjectionModel", "cvssp/audioldm2", "projection_model", {}), + ("StableAudioProjectionModel", "stabilityai/stable-audio-open-1.0", "projection_model", {}), + ("StableAudio3DurationEmbedder", SA3_CHECKPOINT, "duration_embedder", {}), + # black-forest-labs/FLUX.1-Redux-dev is gated: requires accepting the license + `hf auth login`. + ("ReduxImageEncoder", "black-forest-labs/FLUX.1-Redux-dev", "image_embedder", {"torch_dtype": torch.bfloat16}), + ("LTXLatentUpsamplerModel", "Lightricks/ltxv-spatial-upscaler-0.9.7", "latent_upsampler", {}), + ("LTX2LatentUpsamplerModel", "Lightricks/LTX-2", "latent_upsampler", {"torch_dtype": torch.bfloat16}), + ("LTX2TextConnectors", "Lightricks/LTX-2", "connectors", {"torch_dtype": torch.bfloat16}), + ("AceStepConditionEncoder", "ACE-Step/acestep-v15-xl-turbo-diffusers", "condition_encoder", {}), +] + + +# (pipeline_class_name, repo, kwargs). These `model_index.json` files reference the relocated classes by their old +# pipeline folder (e.g. `["ltx2", "LTX2TextConnectors"]`). +PIPELINE_CHECKPOINTS = [ + ("StableDiffusionGLIGENTextImagePipeline", "anhnct/Gligen_Text_Image", {"torch_dtype": torch.float16}), + ("AudioLDM2Pipeline", "cvssp/audioldm2", {"torch_dtype": torch.float16}), + ("StableAudioPipeline", "stabilityai/stable-audio-open-1.0", {"torch_dtype": torch.float16}), + ("StableAudio3Pipeline", SA3_CHECKPOINT, {"torch_dtype": torch.float16}), + ("FluxPriorReduxPipeline", "black-forest-labs/FLUX.1-Redux-dev", {"torch_dtype": torch.bfloat16}), + ("LTXLatentUpsamplePipeline", "Lightricks/ltxv-spatial-upscaler-0.9.7", {"torch_dtype": torch.bfloat16}), + ("LTX2Pipeline", "Lightricks/LTX-2", {"torch_dtype": torch.bfloat16}), + ("AceStepPipeline", "ACE-Step/acestep-v15-xl-turbo-diffusers", {"torch_dtype": torch.bfloat16}), + ( + "SanaWMLTX2Refiner", + "Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refiner", + {"torch_dtype": torch.bfloat16}, + ), +] + + +def _try_load(class_name: str, repo: str, kwargs: dict, subfolder: str | None = None) -> str | None: + """Load the checkpoint; return None on success, or the failure reason.""" + try: + cls = getattr(diffusers, class_name) + load_kwargs = dict(kwargs) + if subfolder is not None: + load_kwargs["subfolder"] = subfolder + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cls.from_pretrained(repo, **load_kwargs) + except Exception: + return traceback.format_exc() + + relocation_warnings = [ + str(w.message) + for w in caught + if issubclass(w.category, FutureWarning) and "diffusers.pipelines." in str(w.message) + ] + if relocation_warnings: + return "Loading emitted relocation deprecation warning(s):\n" + "\n".join(relocation_warnings) + return None + + +def main() -> int: + failures: list[tuple[str, str, str]] = [] # (class_name, target, reason) + + print("=== Model components ===") + for class_name, repo, subfolder, kwargs in MODEL_CHECKPOINTS: + target = repo + (f":{subfolder}" if subfolder else "") + print(f"\n[{class_name}] {target}") + err = _try_load(class_name, repo, kwargs, subfolder) + if err is None: + print(" PASS") + else: + print(" FAIL") + failures.append((class_name, target, err)) + + print("\n=== Pipelines ===") + for class_name, repo, kwargs in PIPELINE_CHECKPOINTS: + print(f"\n[{class_name}] {repo}") + err = _try_load(class_name, repo, kwargs) + if err is None: + print(" PASS") + else: + print(" FAIL") + failures.append((class_name, repo, err)) + + print() + if failures: + print(f"FAILED: {len(failures)} case(s).\n") + for class_name, target, err in failures: + print(f"--- {class_name} ({target}) ---") + print(err) + return 1 + print("OK: all cases passed.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index ac677917242b..3c39878759f0 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -239,6 +239,7 @@ ] _import_structure["models"].extend( [ + "AceStepConditionEncoder", "AceStepTransformer1DModel", "AllegroTransformer3DModel", "AnimaTextConditioner", @@ -246,6 +247,7 @@ "AnyFlowTransformer3DModel", "AsymmetricAutoencoderKL", "AttentionBackendName", + "AudioLDM2ProjectionModel", "AuraFlowTransformer2DModel", "AutoencoderDC", "AutoencoderKL", @@ -283,6 +285,7 @@ "CacheMixin", "ChromaTransformer2DModel", "ChronoEditTransformer3DModel", + "CLIPImageProjection", "CogVideoXTransformer3DModel", "CogView3PlusTransformer2DModel", "CogView4Transformer2DModel", @@ -331,8 +334,11 @@ "LongCatAudioDiTTransformer", "LongCatAudioDiTVae", "LongCatImageTransformer2DModel", + "LTX2LatentUpsamplerModel", + "LTX2TextConnectors", "LTX2VideoDiffusionDecoderModel", "LTX2VideoTransformer3DModel", + "LTXLatentUpsamplerModel", "LTXVideoTransformer3DModel", "Lumina2Transformer2DModel", "LuminaNextDiT2DModel", @@ -360,6 +366,7 @@ "QwenImageControlNetModel", "QwenImageMultiControlNetModel", "QwenImageTransformer2DModel", + "ReduxImageEncoder", "RFD4ProteinaAutoencoder", "RFD4ProteinaTransformer", "SanaControlNetModel", @@ -373,7 +380,9 @@ "SkyReelsV2Transformer3DModel", "SparseControlNetModel", "StableAudio3DiTModel", + "StableAudio3DurationEmbedder", "StableAudioDiTModel", + "StableAudioProjectionModel", "StableCascadeUNet", "T2IAdapter", "T5FilmDecoder", @@ -629,7 +638,6 @@ [ "AceStepAudioTokenDetokenizer", "AceStepAudioTokenizer", - "AceStepConditionEncoder", "AceStepPipeline", "AllegroPipeline", "AltDiffusionImg2ImgPipeline", @@ -647,7 +655,6 @@ "AnyFlowFARPipeline", "AnyFlowPipeline", "AudioLDM2Pipeline", - "AudioLDM2ProjectionModel", "AudioLDM2UNet2DConditionModel", "AudioLDMPipeline", "AuraFlowPipeline", @@ -661,7 +668,6 @@ "ChromaInpaintPipeline", "ChromaPipeline", "ChronoEditPipeline", - "CLIPImageProjection", "CogVideoXFunControlPipeline", "CogVideoXImageToVideoPipeline", "CogVideoXPipeline", @@ -819,7 +825,6 @@ "QwenImageInpaintPipeline", "QwenImageLayeredPipeline", "QwenImagePipeline", - "ReduxImageEncoder", "SanaControlNetPipeline", "SanaImageToVideoPipeline", "SanaPAGPipeline", @@ -840,11 +845,9 @@ "SkyReelsV2ImageToVideoPipeline", "SkyReelsV2Pipeline", "StableAudio3AudioToAudioPipeline", - "StableAudio3DurationEmbedder", "StableAudio3InpaintPipeline", "StableAudio3Pipeline", "StableAudioPipeline", - "StableAudioProjectionModel", "StableCascadeCombinedPipeline", "StableCascadeDecoderPipeline", "StableCascadePriorPipeline", @@ -1161,6 +1164,7 @@ VaeImageProcessorLDM3D, ) from .models import ( + AceStepConditionEncoder, AceStepTransformer1DModel, AllegroTransformer3DModel, AnimaTextConditioner, @@ -1168,6 +1172,7 @@ AnyFlowTransformer3DModel, AsymmetricAutoencoderKL, AttentionBackendName, + AudioLDM2ProjectionModel, AuraFlowTransformer2DModel, AutoencoderDC, AutoencoderKL, @@ -1205,6 +1210,7 @@ CacheMixin, ChromaTransformer2DModel, ChronoEditTransformer3DModel, + CLIPImageProjection, CogVideoXTransformer3DModel, CogView3PlusTransformer2DModel, CogView4Transformer2DModel, @@ -1253,8 +1259,11 @@ LongCatAudioDiTTransformer, LongCatAudioDiTVae, LongCatImageTransformer2DModel, + LTX2LatentUpsamplerModel, + LTX2TextConnectors, LTX2VideoDiffusionDecoderModel, LTX2VideoTransformer3DModel, + LTXLatentUpsamplerModel, LTXVideoTransformer3DModel, Lumina2Transformer2DModel, LuminaNextDiT2DModel, @@ -1282,6 +1291,7 @@ QwenImageControlNetModel, QwenImageMultiControlNetModel, QwenImageTransformer2DModel, + ReduxImageEncoder, RFD4ProteinaAutoencoder, RFD4ProteinaTransformer, SanaControlNetModel, @@ -1295,7 +1305,9 @@ SkyReelsV2Transformer3DModel, SparseControlNetModel, StableAudio3DiTModel, + StableAudio3DurationEmbedder, StableAudioDiTModel, + StableAudioProjectionModel, T2IAdapter, T5FilmDecoder, TensorParallelConfig, @@ -1348,7 +1360,6 @@ AutoPipelineForText2Image, BlipDiffusionControlNetPipeline, BlipDiffusionPipeline, - CLIPImageProjection, ConsistencyModelPipeline, DanceDiffusionPipeline, DDIMPipeline, @@ -1528,7 +1539,6 @@ from .pipelines import ( AceStepAudioTokenDetokenizer, AceStepAudioTokenizer, - AceStepConditionEncoder, AceStepPipeline, AllegroPipeline, AltDiffusionImg2ImgPipeline, @@ -1546,7 +1556,6 @@ AnyFlowFARPipeline, AnyFlowPipeline, AudioLDM2Pipeline, - AudioLDM2ProjectionModel, AudioLDM2UNet2DConditionModel, AudioLDMPipeline, AuraFlowPipeline, @@ -1558,7 +1567,6 @@ ChromaInpaintPipeline, ChromaPipeline, ChronoEditPipeline, - CLIPImageProjection, CogVideoXFunControlPipeline, CogVideoXImageToVideoPipeline, CogVideoXPipeline, @@ -1716,7 +1724,6 @@ QwenImageInpaintPipeline, QwenImageLayeredPipeline, QwenImagePipeline, - ReduxImageEncoder, SanaControlNetPipeline, SanaImageToVideoPipeline, SanaPAGPipeline, @@ -1736,11 +1743,9 @@ SkyReelsV2ImageToVideoPipeline, SkyReelsV2Pipeline, StableAudio3AudioToAudioPipeline, - StableAudio3DurationEmbedder, StableAudio3InpaintPipeline, StableAudio3Pipeline, StableAudioPipeline, - StableAudioProjectionModel, StableCascadeCombinedPipeline, StableCascadeDecoderPipeline, StableCascadePriorPipeline, diff --git a/src/diffusers/loaders/single_file.py b/src/diffusers/loaders/single_file.py index 881ff9b96a4c..91e5320e38e3 100644 --- a/src/diffusers/loaders/single_file.py +++ b/src/diffusers/loaders/single_file.py @@ -64,6 +64,12 @@ def load_single_file_sub_model( disable_mmap=False, **kwargs, ): + from ..pipelines.pipeline_loading_utils import _RELOCATED_PIPELINE_CLASSES + + if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: + library_name = "diffusers" + is_pipeline_module = False + if is_pipeline_module: pipeline_module = getattr(pipelines, library_name) class_obj = getattr(pipeline_module, class_name) diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 3713cdfec22c..c29c178b83c5 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -69,8 +69,15 @@ _import_structure["autoencoders.mmaudio_vocoder"] = ["MMAudioVocoder"] _import_structure["autoencoders.vq_model"] = ["VQModel"] _import_structure["cache_utils"] = ["CacheMixin"] + _import_structure["condition_embedders.condition_embedder_ace_step"] = ["AceStepConditionEncoder"] _import_structure["condition_embedders.condition_embedder_anima"] = ["AnimaTextConditioner"] + _import_structure["condition_embedders.condition_embedder_audioldm2"] = ["AudioLDM2ProjectionModel"] + _import_structure["condition_embedders.condition_embedder_clip_image"] = ["CLIPImageProjection"] + _import_structure["condition_embedders.condition_embedder_flux_redux"] = ["ReduxImageEncoder"] + _import_structure["condition_embedders.condition_embedder_ltx2"] = ["LTX2TextConnectors"] _import_structure["condition_embedders.condition_embedder_minimax_music3"] = ["MiniMaxMusic3ConditionEncoder"] + _import_structure["condition_embedders.condition_embedder_stable_audio"] = ["StableAudioProjectionModel"] + _import_structure["condition_embedders.condition_embedder_stable_audio_3"] = ["StableAudio3DurationEmbedder"] _import_structure["controlnets.controlnet"] = ["ControlNetModel"] _import_structure["controlnets.controlnet_cosmos"] = ["CosmosControlNetModel"] _import_structure["controlnets.controlnet_flux"] = ["FluxControlNetModel", "FluxMultiControlNetModel"] @@ -92,6 +99,8 @@ _import_structure["controlnets.multicontrolnet_union"] = ["MultiControlNetUnionModel"] _import_structure["embeddings"] = ["ImageProjection"] _import_structure["latent_upscaler.latent_upscaler_kandinsky6_sr"] = ["Kandinsky6SRLatentUpscalerBank"] + _import_structure["latent_upscaler.latent_upscaler_ltx"] = ["LTXLatentUpsamplerModel"] + _import_structure["latent_upscaler.latent_upscaler_ltx2"] = ["LTX2LatentUpsamplerModel"] _import_structure["modeling_utils"] = ["ModelMixin"] _import_structure["transformers.ace_step_transformer"] = ["AceStepTransformer1DModel"] _import_structure["transformers.auraflow_transformer_2d"] = ["AuraFlowTransformer2DModel"] @@ -231,7 +240,17 @@ VQModel, ) from .cache_utils import CacheMixin - from .condition_embedders import AnimaTextConditioner, MiniMaxMusic3ConditionEncoder + from .condition_embedders import ( + AceStepConditionEncoder, + AnimaTextConditioner, + AudioLDM2ProjectionModel, + CLIPImageProjection, + LTX2TextConnectors, + MiniMaxMusic3ConditionEncoder, + ReduxImageEncoder, + StableAudio3DurationEmbedder, + StableAudioProjectionModel, + ) from .controlnets import ( ControlNetModel, ControlNetUnionModel, @@ -253,7 +272,7 @@ ZImageControlNetModel, ) from .embeddings import ImageProjection - from .latent_upscaler import Kandinsky6SRLatentUpscalerBank + from .latent_upscaler import Kandinsky6SRLatentUpscalerBank, LTX2LatentUpsamplerModel, LTXLatentUpsamplerModel from .modeling_utils import ModelMixin from .transformers import ( AceStepTransformer1DModel, diff --git a/src/diffusers/models/condition_embedders/__init__.py b/src/diffusers/models/condition_embedders/__init__.py index 2c54cbcde7bc..7ba87e218095 100644 --- a/src/diffusers/models/condition_embedders/__init__.py +++ b/src/diffusers/models/condition_embedders/__init__.py @@ -2,5 +2,12 @@ if is_torch_available(): + from .condition_embedder_ace_step import AceStepConditionEncoder from .condition_embedder_anima import AnimaTextConditioner + from .condition_embedder_audioldm2 import AudioLDM2ProjectionModel + from .condition_embedder_clip_image import CLIPImageProjection + from .condition_embedder_flux_redux import ReduxImageEncoder + from .condition_embedder_ltx2 import LTX2TextConnectors from .condition_embedder_minimax_music3 import MiniMaxMusic3ConditionEncoder + from .condition_embedder_stable_audio import StableAudioProjectionModel + from .condition_embedder_stable_audio_3 import StableAudio3DurationEmbedder diff --git a/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py b/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py new file mode 100644 index 000000000000..d918cc6d02ec --- /dev/null +++ b/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py @@ -0,0 +1,477 @@ +# Copyright 2025 The ACE-Step Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Condition encoder for ACE-Step 1.5. + +Holds the condition encoder (lyric + timbre + text packing), the encoder layer (``AceStepEncoderLayer`` — not used by +the DiT itself, hence kept here), and the ``_pack_sequences`` helper. The DiT uses the RoPE helper, +``AceStepAttention``, and ``_create_4d_mask`` from ``diffusers/models/transformers/ace_step_transformer.py``. +""" + +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin +from ..normalization import RMSNorm +from ..transformers.ace_step_transformer import ( + AceStepAttention, + AceStepMLP, + _ace_step_rotary_freqs, + _create_4d_mask, + _is_flash_attention_backend, +) + + +# --------------------------------------------------------------------------- # +# helpers used only by condition encoder # +# --------------------------------------------------------------------------- # + + +def _pack_sequences( + hidden1: torch.Tensor, hidden2: torch.Tensor, mask1: torch.Tensor, mask2: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Pack two masked sequences into one with all valid tokens first. + + Concatenates ``hidden1`` + ``hidden2`` along the sequence dim, then stably sorts each batch so mask=1 tokens come + before mask=0 tokens. Returns the packed hidden states plus a fresh contiguous mask. + """ + hidden_cat = torch.cat([hidden1, hidden2], dim=1) + mask_cat = torch.cat([mask1, mask2], dim=1) + + B, L, D = hidden_cat.shape + sort_idx = mask_cat.argsort(dim=1, descending=True, stable=True) + hidden_left = torch.gather(hidden_cat, 1, sort_idx.unsqueeze(-1).expand(B, L, D)) + lengths = mask_cat.sum(dim=1) + new_mask = torch.arange(L, dtype=torch.long, device=hidden_cat.device).unsqueeze(0) < lengths.unsqueeze(1) + return hidden_left, new_mask + + +class AceStepEncoderLayer(nn.Module): + """Pre-LN transformer block used by the lyric and timbre encoders.""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + intermediate_size: int, + attention_bias: bool = False, + attention_dropout: float = 0.0, + rms_norm_eps: float = 1e-6, + sliding_window: Optional[int] = None, + ): + super().__init__() + self.self_attn = AceStepAttention( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + bias=attention_bias, + dropout=attention_dropout, + eps=rms_norm_eps, + sliding_window=sliding_window, + is_cross_attention=False, + ) + self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.mlp = AceStepMLP(hidden_size, intermediate_size) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states=hidden_states, + image_rotary_emb=position_embeddings, + attention_mask=attention_mask, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +# --------------------------------------------------------------------------- # +# encoders # +# --------------------------------------------------------------------------- # + + +class AceStepLyricEncoder(ModelMixin, ConfigMixin): + """Lyric encoder: projects Qwen3 lyric embeddings and runs a small transformer. + + Output feeds the DiT cross-attention (after packing with text + timbre). + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + hidden_size: int = 2048, + intermediate_size: int = 6144, + text_hidden_dim: int = 1024, + num_lyric_encoder_hidden_layers: int = 8, + num_attention_heads: int = 16, + num_key_value_heads: int = 8, + head_dim: int = 128, + rope_theta: float = 1000000.0, + attention_bias: bool = False, + attention_dropout: float = 0.0, + rms_norm_eps: float = 1e-6, + sliding_window: int = 128, + layer_types: list = None, + ): + super().__init__() + + if layer_types is None: + layer_types = [ + "sliding_attention" if bool((i + 1) % 2) else "full_attention" + for i in range(num_lyric_encoder_hidden_layers) + ] + + self.embed_tokens = nn.Linear(text_hidden_dim, hidden_size) + self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.head_dim = head_dim + self.rope_theta = rope_theta + self.sliding_window = sliding_window + + self.layers = nn.ModuleList( + [ + AceStepEncoderLayer( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + intermediate_size=intermediate_size, + attention_bias=attention_bias, + attention_dropout=attention_dropout, + rms_norm_eps=rms_norm_eps, + sliding_window=sliding_window if layer_types[i] == "sliding_attention" else None, + ) + for i in range(num_lyric_encoder_hidden_layers) + ] + ) + + self._layer_types = layer_types + self.gradient_checkpointing = False + + def forward( + self, + inputs_embeds: torch.FloatTensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + inputs_embeds = self.embed_tokens(inputs_embeds) + + seq_len = inputs_embeds.shape[1] + dtype = inputs_embeds.dtype + device = inputs_embeds.device + + cos, sin = _ace_step_rotary_freqs(seq_len, self.head_dim, self.rope_theta, device, dtype) + position_embeddings = (cos, sin) + + if _is_flash_attention_backend(self.layers[0].self_attn.processor): + full_attn_mask = attention_mask + sliding_attn_mask = attention_mask + else: + full_attn_mask = _create_4d_mask( + seq_len=seq_len, dtype=dtype, device=device, attention_mask=attention_mask, is_causal=False + ) + sliding_attn_mask = _create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=attention_mask, + sliding_window=self.sliding_window, + is_sliding_window=True, + is_causal=False, + ) + + hidden_states = inputs_embeds + for i, layer_module in enumerate(self.layers): + mask = sliding_attn_mask if self._layer_types[i] == "sliding_attention" else full_attn_mask + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + layer_module, hidden_states, position_embeddings, mask + ) + else: + hidden_states = layer_module( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=mask, + ) + return self.norm(hidden_states) + + +class AceStepTimbreEncoder(ModelMixin, ConfigMixin): + """Timbre encoder: consumes VAE-encoded reference-audio latents and returns a + pooled per-batch timbre embedding (plus a presence mask). + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + hidden_size: int = 2048, + intermediate_size: int = 6144, + timbre_hidden_dim: int = 64, + num_timbre_encoder_hidden_layers: int = 4, + num_attention_heads: int = 16, + num_key_value_heads: int = 8, + head_dim: int = 128, + rope_theta: float = 1000000.0, + attention_bias: bool = False, + attention_dropout: float = 0.0, + rms_norm_eps: float = 1e-6, + sliding_window: int = 128, + layer_types: list = None, + ): + super().__init__() + + if layer_types is None: + layer_types = [ + "sliding_attention" if bool((i + 1) % 2) else "full_attention" + for i in range(num_timbre_encoder_hidden_layers) + ] + + self.embed_tokens = nn.Linear(timbre_hidden_dim, hidden_size) + self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.special_token = nn.Parameter(torch.randn(1, 1, hidden_size)) + self.head_dim = head_dim + self.rope_theta = rope_theta + self.sliding_window = sliding_window + + self.layers = nn.ModuleList( + [ + AceStepEncoderLayer( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + intermediate_size=intermediate_size, + attention_bias=attention_bias, + attention_dropout=attention_dropout, + rms_norm_eps=rms_norm_eps, + sliding_window=sliding_window if layer_types[i] == "sliding_attention" else None, + ) + for i in range(num_timbre_encoder_hidden_layers) + ] + ) + + self._layer_types = layer_types + self.gradient_checkpointing = False + + @staticmethod + def unpack_timbre_embeddings( + timbre_embs_packed: torch.Tensor, refer_audio_order_mask: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + N, d = timbre_embs_packed.shape + device = timbre_embs_packed.device + dtype = timbre_embs_packed.dtype + + B = int(refer_audio_order_mask.max().item() + 1) + counts = torch.bincount(refer_audio_order_mask, minlength=B) + max_count = counts.max().item() + + sorted_indices = torch.argsort(refer_audio_order_mask * N + torch.arange(N, device=device), stable=True) + sorted_batch_ids = refer_audio_order_mask[sorted_indices] + + positions = torch.arange(N, device=device) + batch_starts = torch.cat([torch.tensor([0], device=device), torch.cumsum(counts, dim=0)[:-1]]) + positions_in_sorted = positions - batch_starts[sorted_batch_ids] + + inverse_indices = torch.empty_like(sorted_indices) + inverse_indices[sorted_indices] = torch.arange(N, device=device) + positions_in_batch = positions_in_sorted[inverse_indices] + + indices_2d = refer_audio_order_mask * max_count + positions_in_batch + one_hot = F.one_hot(indices_2d, num_classes=B * max_count).to(dtype) + + timbre_embs_flat = one_hot.t() @ timbre_embs_packed + timbre_embs_unpack = timbre_embs_flat.reshape(B, max_count, d) + + mask_flat = (one_hot.sum(dim=0) > 0).long() + new_mask = mask_flat.reshape(B, max_count) + return timbre_embs_unpack, new_mask + + def forward( + self, + refer_audio_acoustic_hidden_states_packed: torch.FloatTensor, + refer_audio_order_mask: torch.LongTensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + inputs_embeds = self.embed_tokens(refer_audio_acoustic_hidden_states_packed) + + seq_len = inputs_embeds.shape[1] + dtype = inputs_embeds.dtype + device = inputs_embeds.device + + cos, sin = _ace_step_rotary_freqs(seq_len, self.head_dim, self.rope_theta, device, dtype) + position_embeddings = (cos, sin) + + sliding_attn_mask = None + if not _is_flash_attention_backend(self.layers[0].self_attn.processor): + sliding_attn_mask = _create_4d_mask( + seq_len=seq_len, + dtype=dtype, + device=device, + attention_mask=None, + sliding_window=self.sliding_window, + is_sliding_window=True, + is_causal=False, + ) + + hidden_states = inputs_embeds + for i, layer_module in enumerate(self.layers): + # No padding mask on timbre input (pre-packed), so full-attention layers see None. + mask = sliding_attn_mask if self._layer_types[i] == "sliding_attention" else None + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + layer_module, hidden_states, position_embeddings, mask + ) + else: + hidden_states = layer_module( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=mask, + ) + + hidden_states = self.norm(hidden_states) + # CLS-like pooling: first-token embedding per packed sequence. + hidden_states = hidden_states[:, 0, :] + timbre_embs_unpack, timbre_embs_mask = self.unpack_timbre_embeddings(hidden_states, refer_audio_order_mask) + return timbre_embs_unpack, timbre_embs_mask + + +# --------------------------------------------------------------------------- # +# condition encoder # +# --------------------------------------------------------------------------- # + + +class AceStepConditionEncoder(ModelMixin, ConfigMixin): + """Fuses text + lyric + timbre conditioning into the packed sequence used by + the DiT's cross-attention. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + hidden_size: int = 2048, + intermediate_size: int = 6144, + text_hidden_dim: int = 1024, + timbre_hidden_dim: int = 64, + num_lyric_encoder_hidden_layers: int = 8, + num_timbre_encoder_hidden_layers: int = 4, + num_attention_heads: int = 16, + num_key_value_heads: int = 8, + head_dim: int = 128, + rope_theta: float = 1000000.0, + attention_bias: bool = False, + attention_dropout: float = 0.0, + rms_norm_eps: float = 1e-6, + sliding_window: int = 128, + layer_types: list = None, + ): + super().__init__() + + self.text_projector = nn.Linear(text_hidden_dim, hidden_size, bias=False) + + self.lyric_encoder = AceStepLyricEncoder( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + text_hidden_dim=text_hidden_dim, + num_lyric_encoder_hidden_layers=num_lyric_encoder_hidden_layers, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rope_theta=rope_theta, + attention_bias=attention_bias, + attention_dropout=attention_dropout, + rms_norm_eps=rms_norm_eps, + sliding_window=sliding_window, + layer_types=layer_types, + ) + + self.timbre_encoder = AceStepTimbreEncoder( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + timbre_hidden_dim=timbre_hidden_dim, + num_timbre_encoder_hidden_layers=num_timbre_encoder_hidden_layers, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rope_theta=rope_theta, + attention_bias=attention_bias, + attention_dropout=attention_dropout, + rms_norm_eps=rms_norm_eps, + sliding_window=sliding_window, + ) + + # Learned null-condition embedding for classifier-free guidance, trained with + # `cfg_ratio=0.15` in the original model. Broadcast along the sequence dim when used. + self.null_condition_emb = nn.Parameter(torch.randn(1, 1, hidden_size)) + + # Silence latent — VAE-encoded audio-silence, stored as (1, T_long, timbre_hidden_dim). + # When no reference audio is provided, the pipeline slices `silence_latent[:, :timbre_fix_frame, :]` + # and feeds that to the timbre encoder. Passing literal zeros puts the timbre encoder + # OOD and produces drone-like audio (observed on all text2music outputs before this fix). + # The placeholder here is overwritten by the converter with the real encoded silence, + # so its shape just needs to match the timbre-encoder input: last dim is + # `timbre_hidden_dim` (so smaller test configs with `timbre_hidden_dim != 64` also load). + self.register_buffer( + "silence_latent", + torch.zeros(1, 15000, timbre_hidden_dim), + persistent=True, + ) + + def forward( + self, + text_hidden_states: torch.FloatTensor, + text_attention_mask: torch.Tensor, + lyric_hidden_states: torch.FloatTensor, + lyric_attention_mask: torch.Tensor, + refer_audio_acoustic_hidden_states_packed: torch.FloatTensor, + refer_audio_order_mask: torch.LongTensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + text_hidden_states = self.text_projector(text_hidden_states) + + lyric_hidden_states = self.lyric_encoder( + inputs_embeds=lyric_hidden_states, attention_mask=lyric_attention_mask + ) + + timbre_embs_unpack, timbre_embs_mask = self.timbre_encoder( + refer_audio_acoustic_hidden_states_packed, refer_audio_order_mask + ) + + encoder_hidden_states, encoder_attention_mask = _pack_sequences( + lyric_hidden_states, timbre_embs_unpack, lyric_attention_mask, timbre_embs_mask + ) + encoder_hidden_states, encoder_attention_mask = _pack_sequences( + encoder_hidden_states, text_hidden_states, encoder_attention_mask, text_attention_mask + ) + + return encoder_hidden_states, encoder_attention_mask diff --git a/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py b/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py new file mode 100644 index 000000000000..4fb7f62867ba --- /dev/null +++ b/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py @@ -0,0 +1,142 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import BaseOutput +from ..modeling_utils import ModelMixin + + +def add_special_tokens(hidden_states, attention_mask, sos_token, eos_token): + batch_size = hidden_states.shape[0] + + if attention_mask is not None: + # Add two more steps to attn mask + new_attn_mask_step = attention_mask.new_ones((batch_size, 1)) + attention_mask = torch.concat([new_attn_mask_step, attention_mask, new_attn_mask_step], dim=-1) + + # Add the SOS / EOS tokens at the start / end of the sequence respectively + sos_token = sos_token.expand(batch_size, 1, -1) + eos_token = eos_token.expand(batch_size, 1, -1) + hidden_states = torch.concat([sos_token, hidden_states, eos_token], dim=1) + return hidden_states, attention_mask + + +@dataclass +class AudioLDM2ProjectionModelOutput(BaseOutput): + """ + Args: + Class for AudioLDM2 projection layer's outputs. + hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states obtained by linearly projecting the hidden-states for each of the text + encoders and subsequently concatenating them together. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices, formed by concatenating the attention masks + for the two text encoders together. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + """ + + hidden_states: torch.Tensor + attention_mask: torch.LongTensor | None = None + + +class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin): + """ + A simple linear projection model to map two text embeddings to a shared latent space. It also inserts learned + embedding vectors at the start and end of each text embedding sequence respectively. Each variable appended with + `_1` refers to that corresponding to the second text encoder. Otherwise, it is from the first. + + Args: + text_encoder_dim (`int`): + Dimensionality of the text embeddings from the first text encoder (CLAP). + text_encoder_1_dim (`int`): + Dimensionality of the text embeddings from the second text encoder (T5 or VITS). + langauge_model_dim (`int`): + Dimensionality of the text embeddings from the language model (GPT2). + """ + + @register_to_config + def __init__( + self, + text_encoder_dim, + text_encoder_1_dim, + langauge_model_dim, + use_learned_position_embedding=None, + max_seq_length=None, + ): + super().__init__() + # additional projection layers for each text encoder + self.projection = nn.Linear(text_encoder_dim, langauge_model_dim) + self.projection_1 = nn.Linear(text_encoder_1_dim, langauge_model_dim) + + # learnable SOS / EOS token embeddings for each text encoder + self.sos_embed = nn.Parameter(torch.ones(langauge_model_dim)) + self.eos_embed = nn.Parameter(torch.ones(langauge_model_dim)) + + self.sos_embed_1 = nn.Parameter(torch.ones(langauge_model_dim)) + self.eos_embed_1 = nn.Parameter(torch.ones(langauge_model_dim)) + + self.use_learned_position_embedding = use_learned_position_embedding + + # learable positional embedding for vits encoder + if self.use_learned_position_embedding is not None: + self.learnable_positional_embedding = torch.nn.Parameter( + torch.zeros((1, text_encoder_1_dim, max_seq_length)) + ) + + def forward( + self, + hidden_states: torch.Tensor | None = None, + hidden_states_1: torch.Tensor | None = None, + attention_mask: torch.LongTensor | None = None, + attention_mask_1: torch.LongTensor | None = None, + ): + hidden_states = self.projection(hidden_states) + hidden_states, attention_mask = add_special_tokens( + hidden_states, attention_mask, sos_token=self.sos_embed, eos_token=self.eos_embed + ) + + # Add positional embedding for Vits hidden state + if self.use_learned_position_embedding is not None: + hidden_states_1 = (hidden_states_1.permute(0, 2, 1) + self.learnable_positional_embedding).permute(0, 2, 1) + + hidden_states_1 = self.projection_1(hidden_states_1) + hidden_states_1, attention_mask_1 = add_special_tokens( + hidden_states_1, attention_mask_1, sos_token=self.sos_embed_1, eos_token=self.eos_embed_1 + ) + + # concatenate clap and t5 text encoding + hidden_states = torch.cat([hidden_states, hidden_states_1], dim=1) + + # concatenate attention masks + if attention_mask is None and attention_mask_1 is not None: + attention_mask = attention_mask_1.new_ones((hidden_states[:2])) + elif attention_mask is not None and attention_mask_1 is None: + attention_mask_1 = attention_mask.new_ones((hidden_states_1[:2])) + + if attention_mask is not None and attention_mask_1 is not None: + attention_mask = torch.cat([attention_mask, attention_mask_1], dim=-1) + else: + attention_mask = None + + return AudioLDM2ProjectionModelOutput( + hidden_states=hidden_states, + attention_mask=attention_mask, + ) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py b/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py new file mode 100644 index 000000000000..a373f69fa39f --- /dev/null +++ b/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py @@ -0,0 +1,29 @@ +# Copyright 2025 The GLIGEN Authors and HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from torch import nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin + + +class CLIPImageProjection(ModelMixin, ConfigMixin): + @register_to_config + def __init__(self, hidden_size: int = 768): + super().__init__() + self.hidden_size = hidden_size + self.project = nn.Linear(self.hidden_size, self.hidden_size, bias=False) + + def forward(self, x): + return self.project(x) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py b/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py new file mode 100644 index 000000000000..526a8ddba99c --- /dev/null +++ b/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py @@ -0,0 +1,46 @@ +# Copyright 2025 Black Forest Labs and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import BaseOutput +from ..modeling_utils import ModelMixin + + +@dataclass +class ReduxImageEncoderOutput(BaseOutput): + image_embeds: torch.Tensor | None = None + + +class ReduxImageEncoder(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + redux_dim: int = 1152, + txt_in_features: int = 4096, + ) -> None: + super().__init__() + + self.redux_up = nn.Linear(redux_dim, txt_in_features * 3) + self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features) + + def forward(self, x: torch.Tensor) -> ReduxImageEncoderOutput: + projected_x = self.redux_down(nn.functional.silu(self.redux_up(x))) + + return ReduxImageEncoderOutput(image_embeds=projected_x) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py b/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py new file mode 100644 index 000000000000..c119b6a0f6d0 --- /dev/null +++ b/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py @@ -0,0 +1,478 @@ +import math + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ..attention import FeedForward +from ..modeling_utils import ModelMixin +from ..transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoAttnProcessor + + +def per_layer_masked_mean_norm( + text_hidden_states: torch.Tensor, + sequence_lengths: torch.Tensor, + device: str | torch.device, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, +): + """ + Performs per-batch per-layer normalization using a masked mean and range on per-layer text encoder hidden_states. + Respects the padding of the hidden states. + + Args: + text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`): + Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`). + sequence_lengths (`torch.Tensor of shape `(batch_size,)`): + The number of valid (non-padded) tokens for each batch instance. + device: (`str` or `torch.device`, *optional*): + torch device to place the resulting embeddings on + padding_side: (`str`, *optional*, defaults to `"left"`): + Whether the text tokenizer performs padding on the `"left"` or `"right"`. + scale_factor (`int`, *optional*, defaults to `8`): + Scaling factor to multiply the normalized hidden states by. + eps (`float`, *optional*, defaults to `1e-6`): + A small positive value for numerical stability when performing normalization. + + Returns: + `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`: + Normed and flattened text encoder hidden states. + """ + batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape + original_dtype = text_hidden_states.dtype + + # Create padding mask + token_indices = torch.arange(seq_len, device=device).unsqueeze(0) + if padding_side == "right": + # For right padding, valid tokens are from 0 to sequence_length-1 + mask = token_indices < sequence_lengths[:, None] # [batch_size, seq_len] + elif padding_side == "left": + # For left padding, valid tokens are from (T - sequence_length) to T-1 + start_indices = seq_len - sequence_lengths[:, None] # [batch_size, 1] + mask = token_indices >= start_indices # [B, T] + else: + raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") + mask = mask[:, :, None, None] # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1] + + # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len) + masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) + num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) + masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) + + # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len) + x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) + x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) + + # Normalization + normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) + normalized_hidden_states = normalized_hidden_states * scale_factor + + # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers) + normalized_hidden_states = normalized_hidden_states.flatten(2) + mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers) + normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) + normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype) + return normalized_hidden_states + + +def per_token_rms_norm(text_encoder_hidden_states: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + variance = torch.mean(text_encoder_hidden_states**2, dim=2, keepdim=True) + norm_text_encoder_hidden_states = text_encoder_hidden_states * torch.rsqrt(variance + eps) + return norm_text_encoder_hidden_states + + +class LTX2RotaryPosEmbed1d(nn.Module): + """ + 1D rotary positional embeddings (RoPE) for the LTX 2.0 text encoder connectors. + """ + + def __init__( + self, + dim: int, + base_seq_len: int = 4096, + theta: float = 10000.0, + double_precision: bool = True, + rope_type: str = "interleaved", + num_attention_heads: int = 32, + ): + super().__init__() + if rope_type not in ["interleaved", "split"]: + raise ValueError(f"{rope_type=} not supported. Choose between 'interleaved' and 'split'.") + + self.dim = dim + self.base_seq_len = base_seq_len + self.theta = theta + self.double_precision = double_precision + self.rope_type = rope_type + self.num_attention_heads = num_attention_heads + + def forward( + self, + batch_size: int, + pos: int, + device: str | torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + # 1. Get 1D position ids + grid_1d = torch.arange(pos, dtype=torch.float32, device=device) + # Get fractional indices relative to self.base_seq_len + grid_1d = grid_1d / self.base_seq_len + grid = grid_1d.unsqueeze(0).repeat(batch_size, 1) # [batch_size, seq_len] + + # 2. Calculate 1D RoPE frequencies + num_rope_elems = 2 # 1 (because 1D) * 2 (for cos, sin) = 2 + freqs_dtype = torch.float64 if self.double_precision else torch.float32 + pow_indices = torch.pow( + self.theta, + torch.linspace(start=0.0, end=1.0, steps=self.dim // num_rope_elems, dtype=freqs_dtype, device=device), + ) + freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) + + # 3. Matrix-vector outer product between pos ids of shape (batch_size, seq_len) and freqs vector of shape + # (self.dim // 2,). + freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs # [B, seq_len, self.dim // 2] + + # 4. Get real, interleaved (cos, sin) frequencies, padded to self.dim + if self.rope_type == "interleaved": + cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) + sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) + + if self.dim % num_rope_elems != 0: + cos_padding = torch.ones_like(cos_freqs[:, :, : self.dim % num_rope_elems]) + sin_padding = torch.zeros_like(sin_freqs[:, :, : self.dim % num_rope_elems]) + cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) + sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) + + elif self.rope_type == "split": + expected_freqs = self.dim // 2 + current_freqs = freqs.shape[-1] + pad_size = expected_freqs - current_freqs + cos_freq = freqs.cos() + sin_freq = freqs.sin() + + if pad_size != 0: + cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) + sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) + + cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) + sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) + + # Reshape freqs to be compatible with multi-head attention + b = cos_freq.shape[0] + t = cos_freq.shape[1] + + cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1) + sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1) + + cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) + sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) + + return cos_freqs, sin_freqs + + +class LTX2TransformerBlock1d(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + activation_fn: str = "gelu-approximate", + eps: float = 1e-6, + rope_type: str = "interleaved", + apply_gated_attention: bool = False, + ): + super().__init__() + + self.norm1 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) + self.attn1 = LTX2Attention( + query_dim=dim, + heads=num_attention_heads, + kv_heads=num_attention_heads, + dim_head=attention_head_dim, + rope_type=rope_type, + apply_gated_attention=apply_gated_attention, + processor=LTX2AudioVideoAttnProcessor(), + ) + + self.norm2 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) + self.ff = FeedForward(dim, activation_fn=activation_fn) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + rotary_emb: torch.Tensor | None = None, + ) -> torch.Tensor: + norm_hidden_states = self.norm1(hidden_states) + attn_hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, query_rotary_emb=rotary_emb) + hidden_states = hidden_states + attn_hidden_states + + norm_hidden_states = self.norm2(hidden_states) + ff_hidden_states = self.ff(norm_hidden_states) + hidden_states = hidden_states + ff_hidden_states + + return hidden_states + + +class LTX2ConnectorTransformer1d(nn.Module): + """ + A 1D sequence transformer for modalities such as text. + + In LTX 2.0, this is used to process the text encoder hidden states for each of the video and audio streams. + """ + + _supports_gradient_checkpointing = True + + def __init__( + self, + num_attention_heads: int = 30, + attention_head_dim: int = 128, + num_layers: int = 2, + num_learnable_registers: int | None = 128, + rope_base_seq_len: int = 4096, + rope_theta: float = 10000.0, + rope_double_precision: bool = True, + eps: float = 1e-6, + causal_temporal_positioning: bool = False, + rope_type: str = "interleaved", + gated_attention: bool = False, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.inner_dim = num_attention_heads * attention_head_dim + self.causal_temporal_positioning = causal_temporal_positioning + + self.num_learnable_registers = num_learnable_registers + self.learnable_registers = None + if num_learnable_registers is not None: + init_registers = torch.rand(num_learnable_registers, self.inner_dim) * 2.0 - 1.0 + self.learnable_registers = torch.nn.Parameter(init_registers) + + self.rope = LTX2RotaryPosEmbed1d( + self.inner_dim, + base_seq_len=rope_base_seq_len, + theta=rope_theta, + double_precision=rope_double_precision, + rope_type=rope_type, + num_attention_heads=num_attention_heads, + ) + + self.transformer_blocks = torch.nn.ModuleList( + [ + LTX2TransformerBlock1d( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + rope_type=rope_type, + apply_gated_attention=gated_attention, + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = torch.nn.RMSNorm(self.inner_dim, eps=eps, elementwise_affine=False) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + attn_mask_binarize_threshold: float = -9000.0, + ) -> tuple[torch.Tensor, torch.Tensor]: + # hidden_states shape: [batch_size, seq_len, hidden_dim] + # attention_mask shape: [batch_size, seq_len] or [batch_size, 1, 1, seq_len] + batch_size, seq_len, _ = hidden_states.shape + + # 1. Replace padding with learned registers, if using + if self.learnable_registers is not None: + if seq_len % self.num_learnable_registers != 0: + raise ValueError( + f"The `hidden_states` sequence length {hidden_states.shape[1]} should be divisible by the number" + f" of learnable registers {self.num_learnable_registers}" + ) + + num_register_repeats = seq_len // self.num_learnable_registers + registers = ( + self.learnable_registers.unsqueeze(0).expand(num_register_repeats, -1, -1).reshape(seq_len, -1) + ) # [seq_len, inner_dim] + + binary_attn_mask = (attention_mask >= attn_mask_binarize_threshold).int() + if binary_attn_mask.ndim == 4: + binary_attn_mask = binary_attn_mask.squeeze(1).squeeze(1) # [B, 1, 1, L] --> [B, L] + + # Move the valid tokens to the front in their original order and fill the tail + # with registers indexed by absolute position, matching the original LTX + # implementation (`_replace_padded_with_learnable_registers`). A stable argsort + # of the inverted mask gathers valid tokens first while preserving their order. + order = torch.argsort(1 - binary_attn_mask, dim=1, stable=True) # [B, L] + front_aligned = torch.gather(hidden_states, 1, order.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1])) + num_valid = binary_attn_mask.sum(dim=1, keepdim=True) # [B, 1] + positions = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0) # [1, L] + front_mask = (positions < num_valid).unsqueeze(-1) # [B, L, 1] + registers_expanded = registers.unsqueeze(0).expand(batch_size, -1, -1) # [B, L, D] + hidden_states = torch.where(front_mask, front_aligned, registers_expanded.to(hidden_states.dtype)) + + # Overwrite attention_mask with an all-zeros mask if using registers. + attention_mask = torch.zeros_like(attention_mask) + + # 2. Calculate 1D RoPE positional embeddings + rotary_emb = self.rope(batch_size, seq_len, device=hidden_states.device) + + # 3. Run 1D transformer blocks + for block in self.transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(block, hidden_states, attention_mask, rotary_emb) + else: + hidden_states = block(hidden_states, attention_mask=attention_mask, rotary_emb=rotary_emb) + + hidden_states = self.norm_out(hidden_states) + + return hidden_states, attention_mask + + +class LTX2TextConnectors(ModelMixin, PeftAdapterMixin, ConfigMixin): + """ + Text connector stack used by LTX 2.0 to process the packed text encoder hidden states for both the video and audio + streams. + """ + + @register_to_config + def __init__( + self, + caption_channels: int = 3840, # default Gemma-3-12B text encoder hidden_size + text_proj_in_factor: int = 49, # num_layers + 1 for embedding layer = 48 + 1 for Gemma-3-12B + video_connector_num_attention_heads: int = 30, + video_connector_attention_head_dim: int = 128, + video_connector_num_layers: int = 2, + video_connector_num_learnable_registers: int | None = 128, + video_gated_attn: bool = False, + audio_connector_num_attention_heads: int = 30, + audio_connector_attention_head_dim: int = 128, + audio_connector_num_layers: int = 2, + audio_connector_num_learnable_registers: int | None = 128, + audio_gated_attn: bool = False, + connector_rope_base_seq_len: int = 4096, + rope_theta: float = 10000.0, + rope_double_precision: bool = True, + causal_temporal_positioning: bool = False, + rope_type: str = "interleaved", + per_modality_projections: bool = False, + video_hidden_dim: int = 4096, + audio_hidden_dim: int = 2048, + proj_bias: bool = False, + ): + super().__init__() + text_encoder_dim = caption_channels * text_proj_in_factor + if per_modality_projections: + self.video_text_proj_in = nn.Linear(text_encoder_dim, video_hidden_dim, bias=proj_bias) + self.audio_text_proj_in = nn.Linear(text_encoder_dim, audio_hidden_dim, bias=proj_bias) + else: + self.text_proj_in = nn.Linear(text_encoder_dim, caption_channels, bias=proj_bias) + + self.video_connector = LTX2ConnectorTransformer1d( + num_attention_heads=video_connector_num_attention_heads, + attention_head_dim=video_connector_attention_head_dim, + num_layers=video_connector_num_layers, + num_learnable_registers=video_connector_num_learnable_registers, + rope_base_seq_len=connector_rope_base_seq_len, + rope_theta=rope_theta, + rope_double_precision=rope_double_precision, + causal_temporal_positioning=causal_temporal_positioning, + rope_type=rope_type, + gated_attention=video_gated_attn, + ) + self.audio_connector = LTX2ConnectorTransformer1d( + num_attention_heads=audio_connector_num_attention_heads, + attention_head_dim=audio_connector_attention_head_dim, + num_layers=audio_connector_num_layers, + num_learnable_registers=audio_connector_num_learnable_registers, + rope_base_seq_len=connector_rope_base_seq_len, + rope_theta=rope_theta, + rope_double_precision=rope_double_precision, + causal_temporal_positioning=causal_temporal_positioning, + rope_type=rope_type, + gated_attention=audio_gated_attn, + ) + + def forward( + self, + text_encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + padding_side: str = "left", + scale_factor: int = 8, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Given per-layer text encoder hidden_states, extracts features and runs per-modality connectors to get text + embeddings for the LTX-2.X DiT models. + + Args: + text_encoder_hidden_states (`torch.Tensor`)): + Per-layer text encoder hidden_states. Can either be 4D with shape `(batch_size, seq_len, + caption_channels, text_proj_in_factor) or 3D with the last two dimensions flattened. + attention_mask (`torch.Tensor` of shape `(batch_size, seq_len)`): + Multiplicative binary attention mask where 1s indicate unmasked positions and 0s indicate masked + positions. + padding_side (`str`, *optional*, defaults to `"left"`): + The padding side used by the text encoder's text encoder (either `"left"` or `"right"`). Defaults to + `"left"` as this is what the default Gemma3-12B text encoder uses. Only used if + `per_modality_projections` is `False` (LTX-2.0 models). + scale_factor (`int`, *optional*, defaults to `8`): + Scale factor for masked mean/range normalization. Only used if `per_modality_projections` is `False` + (LTX-2.0 models). + """ + if text_encoder_hidden_states.ndim == 3: + # Ensure shape is [batch_size, seq_len, caption_channels, text_proj_in_factor] + text_encoder_hidden_states = text_encoder_hidden_states.unflatten(2, (self.config.caption_channels, -1)) + + if self.config.per_modality_projections: + # LTX-2.3 + norm_text_encoder_hidden_states = per_token_rms_norm(text_encoder_hidden_states) + + norm_text_encoder_hidden_states = norm_text_encoder_hidden_states.flatten(2, 3) + bool_mask = attention_mask.bool().unsqueeze(-1) + norm_text_encoder_hidden_states = torch.where( + bool_mask, norm_text_encoder_hidden_states, torch.zeros_like(norm_text_encoder_hidden_states) + ) + + # Rescale norms with respect to video and audio dims for feature extractors + video_scale_factor = math.sqrt(self.config.video_hidden_dim / self.config.caption_channels) + video_norm_text_emb = norm_text_encoder_hidden_states * video_scale_factor + audio_scale_factor = math.sqrt(self.config.audio_hidden_dim / self.config.caption_channels) + audio_norm_text_emb = norm_text_encoder_hidden_states * audio_scale_factor + + # Per-Modality Feature extractors + video_text_emb_proj = self.video_text_proj_in(video_norm_text_emb) + audio_text_emb_proj = self.audio_text_proj_in(audio_norm_text_emb) + else: + # LTX-2.0 + sequence_lengths = attention_mask.sum(dim=-1) + norm_text_encoder_hidden_states = per_layer_masked_mean_norm( + text_hidden_states=text_encoder_hidden_states, + sequence_lengths=sequence_lengths, + device=text_encoder_hidden_states.device, + padding_side=padding_side, + scale_factor=scale_factor, + ) + + text_emb_proj = self.text_proj_in(norm_text_encoder_hidden_states) + video_text_emb_proj = text_emb_proj + audio_text_emb_proj = text_emb_proj + + # Convert to additive attention mask for connectors + text_dtype = video_text_emb_proj.dtype + attention_mask = (attention_mask.to(torch.int64) - 1).to(text_dtype) + attention_mask = attention_mask.reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) + add_attn_mask = attention_mask * torch.finfo(text_dtype).max + + video_text_embedding, video_attn_mask = self.video_connector(video_text_emb_proj, add_attn_mask) + + # Convert video attn mask to binary (multiplicative) mask and mask video text embedding + binary_attn_mask = (video_attn_mask < 1e-6).to(torch.int64) + binary_attn_mask = binary_attn_mask.reshape(video_text_embedding.shape[0], video_text_embedding.shape[1], 1) + video_text_embedding = video_text_embedding * binary_attn_mask + + audio_text_embedding, _ = self.audio_connector(audio_text_emb_proj, add_attn_mask) + + return video_text_embedding, audio_text_embedding, binary_attn_mask.squeeze(-1) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py b/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py new file mode 100644 index 000000000000..28bf56315939 --- /dev/null +++ b/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py @@ -0,0 +1,156 @@ +# Copyright 2025 Stability AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from math import pi + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import BaseOutput, logging +from ..modeling_utils import ModelMixin + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class StableAudioPositionalEmbedding(nn.Module): + """Used for continuous time""" + + def __init__(self, dim: int): + super().__init__() + assert (dim % 2) == 0 + half_dim = dim // 2 + self.weights = nn.Parameter(torch.randn(half_dim)) + + def forward(self, times: torch.Tensor) -> torch.Tensor: + times = times[..., None] + freqs = times * self.weights[None] * 2 * pi + fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1) + fouriered = torch.cat((times, fouriered), dim=-1) + return fouriered + + +@dataclass +class StableAudioProjectionModelOutput(BaseOutput): + """ + Args: + Class for StableAudio projection layer's outputs. + text_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states obtained by linearly projecting the hidden-states for the text encoder. + seconds_start_hidden_states (`torch.Tensor` of shape `(batch_size, 1, hidden_size)`, *optional*): + Sequence of hidden-states obtained by linearly projecting the audio start hidden states. + seconds_end_hidden_states (`torch.Tensor` of shape `(batch_size, 1, hidden_size)`, *optional*): + Sequence of hidden-states obtained by linearly projecting the audio end hidden states. + """ + + text_hidden_states: torch.Tensor | None = None + seconds_start_hidden_states: torch.Tensor | None = None + seconds_end_hidden_states: torch.Tensor | None = None + + +class StableAudioNumberConditioner(nn.Module): + """ + A simple linear projection model to map numbers to a latent space. + + Args: + number_embedding_dim (`int`): + Dimensionality of the number embeddings. + min_value (`int`): + The minimum value of the seconds number conditioning modules. + max_value (`int`): + The maximum value of the seconds number conditioning modules + internal_dim (`int`): + Dimensionality of the intermediate number hidden states. + """ + + def __init__( + self, + number_embedding_dim, + min_value, + max_value, + internal_dim: int | None = 256, + ): + super().__init__() + self.time_positional_embedding = nn.Sequential( + StableAudioPositionalEmbedding(internal_dim), + nn.Linear(in_features=internal_dim + 1, out_features=number_embedding_dim), + ) + + self.number_embedding_dim = number_embedding_dim + self.min_value = min_value + self.max_value = max_value + + def forward( + self, + floats: torch.Tensor, + ): + floats = floats.clamp(self.min_value, self.max_value) + + normalized_floats = (floats - self.min_value) / (self.max_value - self.min_value) + + # Cast floats to same type as embedder + embedder_dtype = next(self.time_positional_embedding.parameters()).dtype + normalized_floats = normalized_floats.to(embedder_dtype) + + embedding = self.time_positional_embedding(normalized_floats) + float_embeds = embedding.view(-1, 1, self.number_embedding_dim) + + return float_embeds + + +class StableAudioProjectionModel(ModelMixin, ConfigMixin): + """ + A simple linear projection model to map the conditioning values to a shared latent space. + + Args: + text_encoder_dim (`int`): + Dimensionality of the text embeddings from the text encoder (T5). + conditioning_dim (`int`): + Dimensionality of the output conditioning tensors. + min_value (`int`): + The minimum value of the seconds number conditioning modules. + max_value (`int`): + The maximum value of the seconds number conditioning modules + """ + + @register_to_config + def __init__(self, text_encoder_dim, conditioning_dim, min_value, max_value): + super().__init__() + self.text_projection = ( + nn.Identity() if conditioning_dim == text_encoder_dim else nn.Linear(text_encoder_dim, conditioning_dim) + ) + self.start_number_conditioner = StableAudioNumberConditioner(conditioning_dim, min_value, max_value) + self.end_number_conditioner = StableAudioNumberConditioner(conditioning_dim, min_value, max_value) + + def forward( + self, + text_hidden_states: torch.Tensor | None = None, + start_seconds: torch.Tensor | None = None, + end_seconds: torch.Tensor | None = None, + ): + text_hidden_states = ( + text_hidden_states if text_hidden_states is None else self.text_projection(text_hidden_states) + ) + seconds_start_hidden_states = ( + start_seconds if start_seconds is None else self.start_number_conditioner(start_seconds) + ) + seconds_end_hidden_states = end_seconds if end_seconds is None else self.end_number_conditioner(end_seconds) + + return StableAudioProjectionModelOutput( + text_hidden_states=text_hidden_states, + seconds_start_hidden_states=seconds_start_hidden_states, + seconds_end_hidden_states=seconds_end_hidden_states, + ) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_stable_audio_3.py b/src/diffusers/models/condition_embedders/condition_embedder_stable_audio_3.py new file mode 100644 index 000000000000..c314f5a85f75 --- /dev/null +++ b/src/diffusers/models/condition_embedders/condition_embedder_stable_audio_3.py @@ -0,0 +1,109 @@ +# Copyright 2025 Stability AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Auxiliary modules for the Stable Audio 3 pipeline. + +``StableAudio3DurationEmbedder`` mirrors the ``NumberConditioner`` used in the reference implementation of SA3 +(``stable_audio_tools``). It encodes a duration value (in seconds) into a fixed-dimensional vector that is passed to +the DiT as global conditioning (AdaLN). + +Architecture (matches reference ``conditioners.py`` ``NumberConditioner`` with ``fourier_features_type="expo"``): + + 1. Normalize: ``t_norm = clamp(t, min_val, max_val); t_norm = (t_norm - min_val) / (max_val - min_val)`` + 2. ``ExpoFourierFeatures``: fixed exponentially-spaced Fourier basis (ref: + ``stable_audio_tools.models.blocks.ExpoFourierFeatures``) + 3. Linear projection: ``fourier_dim → output_dim`` +""" + +import math + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import logging +from ..modeling_utils import ModelMixin + + +logger = logging.get_logger(__name__) + + +class StableAudio3DurationEmbedder(ModelMixin, ConfigMixin): + """ + Embeds a duration value (in seconds) into a global conditioning vector for the Stable Audio 3 DiT (used as the + ``global_hidden_states`` AdaLN input). + + Replicates ``NumberConditioner(fourier_features_type="expo")`` from the SA3 reference implementation. + + Args: + output_dim (`int`, defaults to 768): + Dimension of the output embedding. Must match the DiT's ``global_cond_dim``. + fourier_dim (`int`, defaults to 256): + Internal Fourier feature dimension (must be even). + min_val (`float`, defaults to 0.0): + Minimum duration value for normalization clamping. + max_val (`float`, defaults to 384.0): + Maximum duration value for normalization clamping. Values above this are clamped. 384 seconds is the + production SA3 Medium upper bound for the ``seconds_total`` conditioner. + min_freq (`float`, defaults to 0.5): + Minimum frequency for the exponential Fourier basis. + max_freq (`float`, defaults to 10000.0): + Maximum frequency for the exponential Fourier basis. + """ + + @register_to_config + def __init__( + self, + output_dim: int = 768, + fourier_dim: int = 256, + min_val: float = 0.0, + max_val: float = 384.0, + min_freq: float = 0.5, + max_freq: float = 10000.0, + ) -> None: + super().__init__() + + # ExpoFourierFeatures — fixed (no learnable parameters) + # Frequencies are exponentially spaced from min_freq to max_freq. + half = fourier_dim // 2 + ramp = torch.linspace(0.0, 1.0, half) + log_min = math.log(min_freq) + log_max = math.log(max_freq) + freqs = torch.exp(ramp * (log_max - log_min) + log_min) + self.register_buffer("freqs", freqs) + + # Linear projection: fourier_dim → output_dim + self.linear = nn.Linear(fourier_dim, output_dim) + + def forward(self, seconds: torch.Tensor) -> torch.Tensor: + """ + Args: + seconds (`torch.Tensor` of shape `(batch,)`): + Duration values in seconds. + + Returns: + `torch.Tensor` of shape `(batch, output_dim)` — duration embeddings. + """ + # 1. Normalize to [0, 1] + seconds = seconds.float().clamp(self.config.min_val, self.config.max_val) + t_norm = (seconds - self.config.min_val) / (self.config.max_val - self.config.min_val) + + # 2. Exponential Fourier features — run in fp32 for stability + t_norm = t_norm.reshape(-1, 1) # (B, 1) + args = t_norm * self.freqs.unsqueeze(0) * 2.0 * math.pi # (B, half) + fourier = torch.cat([args.cos(), args.sin()], dim=-1) # (B, fourier_dim) + + # 3. Linear projection + return self.linear(fourier.to(self.dtype)) # (B, output_dim) diff --git a/src/diffusers/models/latent_upscaler/__init__.py b/src/diffusers/models/latent_upscaler/__init__.py index 1c3f83192f85..f48ab36563be 100644 --- a/src/diffusers/models/latent_upscaler/__init__.py +++ b/src/diffusers/models/latent_upscaler/__init__.py @@ -3,3 +3,5 @@ if is_torch_available(): from .latent_upscaler_kandinsky6_sr import Kandinsky6SRLatentUpscalerBank + from .latent_upscaler_ltx import LTXLatentUpsamplerModel + from .latent_upscaler_ltx2 import LTX2LatentUpsamplerModel diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py new file mode 100644 index 000000000000..76bb35ead842 --- /dev/null +++ b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py @@ -0,0 +1,186 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin + + +class ResBlock(torch.nn.Module): + def __init__(self, channels: int, mid_channels: int | None = None, dims: int = 3): + super().__init__() + if mid_channels is None: + mid_channels = channels + + Conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + + self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1) + self.norm1 = torch.nn.GroupNorm(32, mid_channels) + self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1) + self.norm2 = torch.nn.GroupNorm(32, channels) + self.activation = torch.nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + hidden_states = self.conv1(hidden_states) + hidden_states = self.norm1(hidden_states) + hidden_states = self.activation(hidden_states) + hidden_states = self.conv2(hidden_states) + hidden_states = self.norm2(hidden_states) + hidden_states = self.activation(hidden_states + residual) + return hidden_states + + +class PixelShuffleND(torch.nn.Module): + def __init__(self, dims, upscale_factors=(2, 2, 2)): + super().__init__() + + self.dims = dims + self.upscale_factors = upscale_factors + + if dims not in [1, 2, 3]: + raise ValueError("dims must be 1, 2, or 3") + + def forward(self, x): + if self.dims == 3: + # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3) + return ( + x.unflatten(1, (-1, *self.upscale_factors[:3])) + .permute(0, 1, 5, 2, 6, 3, 7, 4) + .flatten(6, 7) + .flatten(4, 5) + .flatten(2, 3) + ) + elif self.dims == 2: + # spatial: b (c p1 p2) h w -> b c (h p1) (w p2) + return ( + x.unflatten(1, (-1, *self.upscale_factors[:2])).permute(0, 1, 4, 2, 5, 3).flatten(4, 5).flatten(2, 3) + ) + elif self.dims == 1: + # temporal: b (c p1) f h w -> b c (f p1) h w + return x.unflatten(1, (-1, *self.upscale_factors[:1])).permute(0, 1, 3, 2, 4, 5).flatten(2, 3) + + +class LTXLatentUpsamplerModel(ModelMixin, ConfigMixin): + """ + Model to spatially upsample VAE latents. + + Args: + in_channels (`int`, defaults to `128`): + Number of channels in the input latent + mid_channels (`int`, defaults to `512`): + Number of channels in the middle layers + num_blocks_per_stage (`int`, defaults to `4`): + Number of ResBlocks to use in each stage (pre/post upsampling) + dims (`int`, defaults to `3`): + Number of dimensions for convolutions (2 or 3) + spatial_upsample (`bool`, defaults to `True`): + Whether to spatially upsample the latent + temporal_upsample (`bool`, defaults to `False`): + Whether to temporally upsample the latent + """ + + @register_to_config + def __init__( + self, + in_channels: int = 128, + mid_channels: int = 512, + num_blocks_per_stage: int = 4, + dims: int = 3, + spatial_upsample: bool = True, + temporal_upsample: bool = False, + ): + super().__init__() + + self.in_channels = in_channels + self.mid_channels = mid_channels + self.num_blocks_per_stage = num_blocks_per_stage + self.dims = dims + self.spatial_upsample = spatial_upsample + self.temporal_upsample = temporal_upsample + + ConvNd = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + + self.initial_conv = ConvNd(in_channels, mid_channels, kernel_size=3, padding=1) + self.initial_norm = torch.nn.GroupNorm(32, mid_channels) + self.initial_activation = torch.nn.SiLU() + + self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) + + if spatial_upsample and temporal_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(3), + ) + elif spatial_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(2), + ) + elif temporal_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(1), + ) + else: + raise ValueError("Either spatial_upsample or temporal_upsample must be True") + + self.post_upsample_res_blocks = torch.nn.ModuleList( + [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + ) + + self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + + if self.dims == 2: + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) + hidden_states = self.initial_conv(hidden_states) + hidden_states = self.initial_norm(hidden_states) + hidden_states = self.initial_activation(hidden_states) + + for block in self.res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.upsampler(hidden_states) + + for block in self.post_upsample_res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.final_conv(hidden_states) + hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) + else: + hidden_states = self.initial_conv(hidden_states) + hidden_states = self.initial_norm(hidden_states) + hidden_states = self.initial_activation(hidden_states) + + for block in self.res_blocks: + hidden_states = block(hidden_states) + + if self.temporal_upsample: + hidden_states = self.upsampler(hidden_states) + hidden_states = hidden_states[:, :, 1:, :, :] + else: + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) + hidden_states = self.upsampler(hidden_states) + hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) + + for block in self.post_upsample_res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.final_conv(hidden_states) + + return hidden_states diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py new file mode 100644 index 000000000000..d43a77920166 --- /dev/null +++ b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py @@ -0,0 +1,285 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin + + +RATIONAL_RESAMPLER_SCALE_MAPPING = { + 0.75: (3, 4), + 1.5: (3, 2), + 2.0: (2, 1), + 4.0: (4, 1), +} + + +# Copied from diffusers.models.latent_upscaler.latent_upscaler_ltx.ResBlock +class ResBlock(torch.nn.Module): + def __init__(self, channels: int, mid_channels: int | None = None, dims: int = 3): + super().__init__() + if mid_channels is None: + mid_channels = channels + + Conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + + self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1) + self.norm1 = torch.nn.GroupNorm(32, mid_channels) + self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1) + self.norm2 = torch.nn.GroupNorm(32, channels) + self.activation = torch.nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + hidden_states = self.conv1(hidden_states) + hidden_states = self.norm1(hidden_states) + hidden_states = self.activation(hidden_states) + hidden_states = self.conv2(hidden_states) + hidden_states = self.norm2(hidden_states) + hidden_states = self.activation(hidden_states + residual) + return hidden_states + + +# Copied from diffusers.models.latent_upscaler.latent_upscaler_ltx.PixelShuffleND +class PixelShuffleND(torch.nn.Module): + def __init__(self, dims, upscale_factors=(2, 2, 2)): + super().__init__() + + self.dims = dims + self.upscale_factors = upscale_factors + + if dims not in [1, 2, 3]: + raise ValueError("dims must be 1, 2, or 3") + + def forward(self, x): + if self.dims == 3: + # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3) + return ( + x.unflatten(1, (-1, *self.upscale_factors[:3])) + .permute(0, 1, 5, 2, 6, 3, 7, 4) + .flatten(6, 7) + .flatten(4, 5) + .flatten(2, 3) + ) + elif self.dims == 2: + # spatial: b (c p1 p2) h w -> b c (h p1) (w p2) + return ( + x.unflatten(1, (-1, *self.upscale_factors[:2])).permute(0, 1, 4, 2, 5, 3).flatten(4, 5).flatten(2, 3) + ) + elif self.dims == 1: + # temporal: b (c p1) f h w -> b c (f p1) h w + return x.unflatten(1, (-1, *self.upscale_factors[:1])).permute(0, 1, 3, 2, 4, 5).flatten(2, 3) + + +class BlurDownsample(torch.nn.Module): + """ + Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel. Applies only on H,W. + Works for dims=2 or dims=3 (per-frame). + """ + + def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None: + super().__init__() + + if dims not in (2, 3): + raise ValueError(f"`dims` must be either 2 or 3 but is {dims}") + if kernel_size < 3 or kernel_size % 2 != 1: + raise ValueError(f"`kernel_size` must be an odd number >= 3 but is {kernel_size}") + + self.dims = dims + self.stride = stride + self.kernel_size = kernel_size + + # 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from + # the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and + # provides a smooth approximation of a Gaussian filter (often called a "binomial filter"). + # The 2D kernel is constructed as the outer product and normalized. + k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)]) + k2d = k[:, None] @ k[None, :] + k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size) + self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.stride == 1: + return x + + if self.dims == 2: + c = x.shape[1] + weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise + x = F.conv2d(x, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c) + else: + # dims == 3: apply per-frame on H,W + b, c, f, _, _ = x.shape + x = x.transpose(1, 2).flatten(0, 1) # [B, C, F, H, W] --> [B * F, C, H, W] + + weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise + x = F.conv2d(x, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c) + + h2, w2 = x.shape[-2:] + x = x.unflatten(0, (b, f)).reshape(b, -1, f, h2, w2) # [B * F, C, H, W] --> [B, C, F, H, W] + return x + + +class SpatialRationalResampler(torch.nn.Module): + """ + Scales by the spatial size of the input by a rational number `scale`. For example, `scale = 0.75` will downsample + by a factor of 3 / 4, while `scale = 1.5` will upsample by a factor of 3 / 2. This works by first upsampling the + input by the (integer) numerator of `scale`, and then performing a blur + stride anti-aliased downsample by the + (integer) denominator. + """ + + def __init__(self, mid_channels: int = 1024, scale: float = 2.0): + super().__init__() + self.scale = float(scale) + num_denom = RATIONAL_RESAMPLER_SCALE_MAPPING.get(scale, None) + if num_denom is None: + raise ValueError( + f"The supplied `scale` {scale} is not supported; supported scales are {list(RATIONAL_RESAMPLER_SCALE_MAPPING.keys())}" + ) + self.num, self.den = num_denom + + self.conv = torch.nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1) + self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) + self.blur_down = BlurDownsample(dims=2, stride=self.den) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Expected x shape: [B * F, C, H, W] + # b, _, f, h, w = x.shape + # x = x.transpose(1, 2).flatten(0, 1) # [B, C, F, H, W] --> [B * F, C, H, W] + x = self.conv(x) + x = self.pixel_shuffle(x) + x = self.blur_down(x) + # x = x.unflatten(0, (b, f)).reshape(b, -1, f, h, w) # [B * F, C, H, W] --> [B, C, F, H, W] + return x + + +class LTX2LatentUpsamplerModel(ModelMixin, ConfigMixin): + """ + Model to spatially upsample VAE latents. + + Args: + in_channels (`int`, defaults to `128`): + Number of channels in the input latent + mid_channels (`int`, defaults to `512`): + Number of channels in the middle layers + num_blocks_per_stage (`int`, defaults to `4`): + Number of ResBlocks to use in each stage (pre/post upsampling) + dims (`int`, defaults to `3`): + Number of dimensions for convolutions (2 or 3) + spatial_upsample (`bool`, defaults to `True`): + Whether to spatially upsample the latent + temporal_upsample (`bool`, defaults to `False`): + Whether to temporally upsample the latent + """ + + @register_to_config + def __init__( + self, + in_channels: int = 128, + mid_channels: int = 1024, + num_blocks_per_stage: int = 4, + dims: int = 3, + spatial_upsample: bool = True, + temporal_upsample: bool = False, + rational_spatial_scale: float = 2.0, + use_rational_resampler: bool = True, + ): + super().__init__() + + self.in_channels = in_channels + self.mid_channels = mid_channels + self.num_blocks_per_stage = num_blocks_per_stage + self.dims = dims + self.spatial_upsample = spatial_upsample + self.temporal_upsample = temporal_upsample + + ConvNd = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + + self.initial_conv = ConvNd(in_channels, mid_channels, kernel_size=3, padding=1) + self.initial_norm = torch.nn.GroupNorm(32, mid_channels) + self.initial_activation = torch.nn.SiLU() + + self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) + + if spatial_upsample and temporal_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(3), + ) + elif spatial_upsample: + if use_rational_resampler: + self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=rational_spatial_scale) + else: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(2), + ) + elif temporal_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), + PixelShuffleND(1), + ) + else: + raise ValueError("Either spatial_upsample or temporal_upsample must be True") + + self.post_upsample_res_blocks = torch.nn.ModuleList( + [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + ) + + self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + + if self.dims == 2: + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) + hidden_states = self.initial_conv(hidden_states) + hidden_states = self.initial_norm(hidden_states) + hidden_states = self.initial_activation(hidden_states) + + for block in self.res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.upsampler(hidden_states) + + for block in self.post_upsample_res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.final_conv(hidden_states) + hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) + else: + hidden_states = self.initial_conv(hidden_states) + hidden_states = self.initial_norm(hidden_states) + hidden_states = self.initial_activation(hidden_states) + + for block in self.res_blocks: + hidden_states = block(hidden_states) + + if self.temporal_upsample: + hidden_states = self.upsampler(hidden_states) + hidden_states = hidden_states[:, :, 1:, :, :] + else: + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) + hidden_states = self.upsampler(hidden_states) + hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) + + for block in self.post_upsample_res_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.final_conv(hidden_states) + + return hidden_states diff --git a/src/diffusers/modular_pipelines/echo/encoders.py b/src/diffusers/modular_pipelines/echo/encoders.py index 5dfe33f387c2..fc22eec421fc 100644 --- a/src/diffusers/modular_pipelines/echo/encoders.py +++ b/src/diffusers/modular_pipelines/echo/encoders.py @@ -23,7 +23,7 @@ from transformers import PreTrainedModel, PreTrainedTokenizerBase from ...models import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video -from ...pipelines.ltx2.connectors import LTX2TextConnectors +from ...models.condition_embedders import LTX2TextConnectors from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam diff --git a/src/diffusers/modular_pipelines/ltx2/encoders.py b/src/diffusers/modular_pipelines/ltx2/encoders.py index 433879d46cf4..682cee1e85fe 100644 --- a/src/diffusers/modular_pipelines/ltx2/encoders.py +++ b/src/diffusers/modular_pipelines/ltx2/encoders.py @@ -22,15 +22,14 @@ from ...configuration_utils import FrozenDict from ...models import AutoencoderKLLTX2Video, LTX2VideoTransformer3DModel - -# NOTE (modular.md gotcha #1): `LTX2TextConnectors`, `LTX2DurationHead`, `LTX2VideoCondition`, -# `LTX2ReferenceCondition`, the prompt-enhancement config/helpers, and the system prompts live under -# `diffusers.pipelines.ltx2.*`, and modular blocks must not import from `diffusers.pipelines.*`. -# `LTX2TextConnectors` / `LTX2DurationHead` are `ModelMixin` / `ConfigMixin` model classes (relocate to -# `src/diffusers/models/`); the two condition dataclasses, the enhancement config, the response/image helpers, and -# the system prompts are plain data and utilities (relocate to a neutral shared module or copy into this package). -# Imported from the pipelines path here only so the draft is runnable. -from ...pipelines.ltx2.connectors import LTX2TextConnectors +from ...models.condition_embedders import LTX2TextConnectors + +# NOTE (modular.md gotcha #1): `LTX2DurationHead`, `LTX2VideoCondition`, `LTX2ReferenceCondition`, the +# prompt-enhancement config/helpers, and the system prompts live under `diffusers.pipelines.ltx2.*`, and modular blocks +# must not import from `diffusers.pipelines.*`. `LTX2DurationHead` is a `ModelMixin` / `ConfigMixin` model class +# (relocate to `src/diffusers/models/`); the two condition dataclasses, the enhancement config, the response/image +# helpers, and the system prompts are plain data and utilities (relocate to a neutral shared module or copy into this +# package). Imported from the pipelines path here only so the draft is runnable. from ...pipelines.ltx2.duration_head import LTX2DurationHead from ...pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition from ...pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition diff --git a/src/diffusers/pipelines/ace_step/modeling_ace_step.py b/src/diffusers/pipelines/ace_step/modeling_ace_step.py index 769b07044420..26f6dcee3ffa 100644 --- a/src/diffusers/pipelines/ace_step/modeling_ace_step.py +++ b/src/diffusers/pipelines/ace_step/modeling_ace_step.py @@ -14,10 +14,9 @@ """Pipeline-specific models for ACE-Step 1.5. -Holds the condition encoder (lyric + timbre + text packing), the encoder layer (``AceStepEncoderLayer`` — not used by -the DiT itself, hence kept here), the audio tokenizer / detokenizer used by cover conditioning, and the -``_pack_sequences`` helper. The DiT uses the RoPE helper, ``AceStepAttention``, and ``_create_4d_mask`` from -``diffusers/models/transformers/ace_step_transformer.py``. +Holds the audio tokenizer / detokenizer used by cover conditioning. The condition encoder now lives in +``diffusers/models/condition_embedders/condition_embedder_ace_step.py``; the names below are kept as deprecated aliases +for backward compatibility. """ from typing import Optional, Tuple @@ -27,349 +26,28 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...models.condition_embedders.condition_embedder_ace_step import ( + AceStepConditionEncoder as _AceStepConditionEncoder, +) +from ...models.condition_embedders.condition_embedder_ace_step import ( + AceStepEncoderLayer, + _pack_sequences, # noqa: F401 +) +from ...models.condition_embedders.condition_embedder_ace_step import AceStepLyricEncoder as _AceStepLyricEncoder +from ...models.condition_embedders.condition_embedder_ace_step import AceStepTimbreEncoder as _AceStepTimbreEncoder from ...models.modeling_utils import ModelMixin from ...models.normalization import RMSNorm from ...models.transformers.ace_step_transformer import ( - AceStepAttention, - AceStepMLP, _ace_step_rotary_freqs, _create_4d_mask, _is_flash_attention_backend, ) -from ...utils import logging +from ...utils import deprecate, logging logger = logging.get_logger(__name__) # pylint: disable=invalid-name -# --------------------------------------------------------------------------- # -# helpers used only by condition encoder # -# --------------------------------------------------------------------------- # - - -def _pack_sequences( - hidden1: torch.Tensor, hidden2: torch.Tensor, mask1: torch.Tensor, mask2: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - """Pack two masked sequences into one with all valid tokens first. - - Concatenates ``hidden1`` + ``hidden2`` along the sequence dim, then stably sorts each batch so mask=1 tokens come - before mask=0 tokens. Returns the packed hidden states plus a fresh contiguous mask. - """ - hidden_cat = torch.cat([hidden1, hidden2], dim=1) - mask_cat = torch.cat([mask1, mask2], dim=1) - - B, L, D = hidden_cat.shape - sort_idx = mask_cat.argsort(dim=1, descending=True, stable=True) - hidden_left = torch.gather(hidden_cat, 1, sort_idx.unsqueeze(-1).expand(B, L, D)) - lengths = mask_cat.sum(dim=1) - new_mask = torch.arange(L, dtype=torch.long, device=hidden_cat.device).unsqueeze(0) < lengths.unsqueeze(1) - return hidden_left, new_mask - - -class AceStepEncoderLayer(nn.Module): - """Pre-LN transformer block used by the lyric and timbre encoders.""" - - def __init__( - self, - hidden_size: int, - num_attention_heads: int, - num_key_value_heads: int, - head_dim: int, - intermediate_size: int, - attention_bias: bool = False, - attention_dropout: float = 0.0, - rms_norm_eps: float = 1e-6, - sliding_window: Optional[int] = None, - ): - super().__init__() - self.self_attn = AceStepAttention( - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - head_dim=head_dim, - bias=attention_bias, - dropout=attention_dropout, - eps=rms_norm_eps, - sliding_window=sliding_window, - is_cross_attention=False, - ) - self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.mlp = AceStepMLP(hidden_size, intermediate_size) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], - attention_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn( - hidden_states=hidden_states, - image_rotary_emb=position_embeddings, - attention_mask=attention_mask, - ) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - return hidden_states - - -# --------------------------------------------------------------------------- # -# encoders # -# --------------------------------------------------------------------------- # - - -class AceStepLyricEncoder(ModelMixin, ConfigMixin): - """Lyric encoder: projects Qwen3 lyric embeddings and runs a small transformer. - - Output feeds the DiT cross-attention (after packing with text + timbre). - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - hidden_size: int = 2048, - intermediate_size: int = 6144, - text_hidden_dim: int = 1024, - num_lyric_encoder_hidden_layers: int = 8, - num_attention_heads: int = 16, - num_key_value_heads: int = 8, - head_dim: int = 128, - rope_theta: float = 1000000.0, - attention_bias: bool = False, - attention_dropout: float = 0.0, - rms_norm_eps: float = 1e-6, - sliding_window: int = 128, - layer_types: list = None, - ): - super().__init__() - - if layer_types is None: - layer_types = [ - "sliding_attention" if bool((i + 1) % 2) else "full_attention" - for i in range(num_lyric_encoder_hidden_layers) - ] - - self.embed_tokens = nn.Linear(text_hidden_dim, hidden_size) - self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.head_dim = head_dim - self.rope_theta = rope_theta - self.sliding_window = sliding_window - - self.layers = nn.ModuleList( - [ - AceStepEncoderLayer( - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - head_dim=head_dim, - intermediate_size=intermediate_size, - attention_bias=attention_bias, - attention_dropout=attention_dropout, - rms_norm_eps=rms_norm_eps, - sliding_window=sliding_window if layer_types[i] == "sliding_attention" else None, - ) - for i in range(num_lyric_encoder_hidden_layers) - ] - ) - - self._layer_types = layer_types - self.gradient_checkpointing = False - - def forward( - self, - inputs_embeds: torch.FloatTensor, - attention_mask: torch.Tensor, - ) -> torch.Tensor: - inputs_embeds = self.embed_tokens(inputs_embeds) - - seq_len = inputs_embeds.shape[1] - dtype = inputs_embeds.dtype - device = inputs_embeds.device - - cos, sin = _ace_step_rotary_freqs(seq_len, self.head_dim, self.rope_theta, device, dtype) - position_embeddings = (cos, sin) - - if _is_flash_attention_backend(self.layers[0].self_attn.processor): - full_attn_mask = attention_mask - sliding_attn_mask = attention_mask - else: - full_attn_mask = _create_4d_mask( - seq_len=seq_len, dtype=dtype, device=device, attention_mask=attention_mask, is_causal=False - ) - sliding_attn_mask = _create_4d_mask( - seq_len=seq_len, - dtype=dtype, - device=device, - attention_mask=attention_mask, - sliding_window=self.sliding_window, - is_sliding_window=True, - is_causal=False, - ) - - hidden_states = inputs_embeds - for i, layer_module in enumerate(self.layers): - mask = sliding_attn_mask if self._layer_types[i] == "sliding_attention" else full_attn_mask - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - layer_module, hidden_states, position_embeddings, mask - ) - else: - hidden_states = layer_module( - hidden_states=hidden_states, - position_embeddings=position_embeddings, - attention_mask=mask, - ) - return self.norm(hidden_states) - - -class AceStepTimbreEncoder(ModelMixin, ConfigMixin): - """Timbre encoder: consumes VAE-encoded reference-audio latents and returns a - pooled per-batch timbre embedding (plus a presence mask). - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - hidden_size: int = 2048, - intermediate_size: int = 6144, - timbre_hidden_dim: int = 64, - num_timbre_encoder_hidden_layers: int = 4, - num_attention_heads: int = 16, - num_key_value_heads: int = 8, - head_dim: int = 128, - rope_theta: float = 1000000.0, - attention_bias: bool = False, - attention_dropout: float = 0.0, - rms_norm_eps: float = 1e-6, - sliding_window: int = 128, - layer_types: list = None, - ): - super().__init__() - - if layer_types is None: - layer_types = [ - "sliding_attention" if bool((i + 1) % 2) else "full_attention" - for i in range(num_timbre_encoder_hidden_layers) - ] - - self.embed_tokens = nn.Linear(timbre_hidden_dim, hidden_size) - self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.special_token = nn.Parameter(torch.randn(1, 1, hidden_size)) - self.head_dim = head_dim - self.rope_theta = rope_theta - self.sliding_window = sliding_window - - self.layers = nn.ModuleList( - [ - AceStepEncoderLayer( - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - head_dim=head_dim, - intermediate_size=intermediate_size, - attention_bias=attention_bias, - attention_dropout=attention_dropout, - rms_norm_eps=rms_norm_eps, - sliding_window=sliding_window if layer_types[i] == "sliding_attention" else None, - ) - for i in range(num_timbre_encoder_hidden_layers) - ] - ) - - self._layer_types = layer_types - self.gradient_checkpointing = False - - @staticmethod - def unpack_timbre_embeddings( - timbre_embs_packed: torch.Tensor, refer_audio_order_mask: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: - N, d = timbre_embs_packed.shape - device = timbre_embs_packed.device - dtype = timbre_embs_packed.dtype - - B = int(refer_audio_order_mask.max().item() + 1) - counts = torch.bincount(refer_audio_order_mask, minlength=B) - max_count = counts.max().item() - - sorted_indices = torch.argsort(refer_audio_order_mask * N + torch.arange(N, device=device), stable=True) - sorted_batch_ids = refer_audio_order_mask[sorted_indices] - - positions = torch.arange(N, device=device) - batch_starts = torch.cat([torch.tensor([0], device=device), torch.cumsum(counts, dim=0)[:-1]]) - positions_in_sorted = positions - batch_starts[sorted_batch_ids] - - inverse_indices = torch.empty_like(sorted_indices) - inverse_indices[sorted_indices] = torch.arange(N, device=device) - positions_in_batch = positions_in_sorted[inverse_indices] - - indices_2d = refer_audio_order_mask * max_count + positions_in_batch - one_hot = F.one_hot(indices_2d, num_classes=B * max_count).to(dtype) - - timbre_embs_flat = one_hot.t() @ timbre_embs_packed - timbre_embs_unpack = timbre_embs_flat.reshape(B, max_count, d) - - mask_flat = (one_hot.sum(dim=0) > 0).long() - new_mask = mask_flat.reshape(B, max_count) - return timbre_embs_unpack, new_mask - - def forward( - self, - refer_audio_acoustic_hidden_states_packed: torch.FloatTensor, - refer_audio_order_mask: torch.LongTensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: - inputs_embeds = self.embed_tokens(refer_audio_acoustic_hidden_states_packed) - - seq_len = inputs_embeds.shape[1] - dtype = inputs_embeds.dtype - device = inputs_embeds.device - - cos, sin = _ace_step_rotary_freqs(seq_len, self.head_dim, self.rope_theta, device, dtype) - position_embeddings = (cos, sin) - - sliding_attn_mask = None - if not _is_flash_attention_backend(self.layers[0].self_attn.processor): - sliding_attn_mask = _create_4d_mask( - seq_len=seq_len, - dtype=dtype, - device=device, - attention_mask=None, - sliding_window=self.sliding_window, - is_sliding_window=True, - is_causal=False, - ) - - hidden_states = inputs_embeds - for i, layer_module in enumerate(self.layers): - # No padding mask on timbre input (pre-packed), so full-attention layers see None. - mask = sliding_attn_mask if self._layer_types[i] == "sliding_attention" else None - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - layer_module, hidden_states, position_embeddings, mask - ) - else: - hidden_states = layer_module( - hidden_states=hidden_states, - position_embeddings=position_embeddings, - attention_mask=mask, - ) - - hidden_states = self.norm(hidden_states) - # CLS-like pooling: first-token embedding per packed sequence. - hidden_states = hidden_states[:, 0, :] - timbre_embs_unpack, timbre_embs_mask = self.unpack_timbre_embeddings(hidden_states, refer_audio_order_mask) - return timbre_embs_unpack, timbre_embs_mask - - # --------------------------------------------------------------------------- # # audio tokenizer / detokenizer # # --------------------------------------------------------------------------- # @@ -744,113 +422,36 @@ def tokenize( return self(hidden_states) -# --------------------------------------------------------------------------- # -# condition encoder # -# --------------------------------------------------------------------------- # - - -class AceStepConditionEncoder(ModelMixin, ConfigMixin): - """Fuses text + lyric + timbre conditioning into the packed sequence used by - the DiT's cross-attention. - """ - - _supports_gradient_checkpointing = True - - @register_to_config - def __init__( - self, - hidden_size: int = 2048, - intermediate_size: int = 6144, - text_hidden_dim: int = 1024, - timbre_hidden_dim: int = 64, - num_lyric_encoder_hidden_layers: int = 8, - num_timbre_encoder_hidden_layers: int = 4, - num_attention_heads: int = 16, - num_key_value_heads: int = 8, - head_dim: int = 128, - rope_theta: float = 1000000.0, - attention_bias: bool = False, - attention_dropout: float = 0.0, - rms_norm_eps: float = 1e-6, - sliding_window: int = 128, - layer_types: list = None, - ): - super().__init__() - - self.text_projector = nn.Linear(text_hidden_dim, hidden_size, bias=False) - - self.lyric_encoder = AceStepLyricEncoder( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - text_hidden_dim=text_hidden_dim, - num_lyric_encoder_hidden_layers=num_lyric_encoder_hidden_layers, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - head_dim=head_dim, - rope_theta=rope_theta, - attention_bias=attention_bias, - attention_dropout=attention_dropout, - rms_norm_eps=rms_norm_eps, - sliding_window=sliding_window, - layer_types=layer_types, - ) - - self.timbre_encoder = AceStepTimbreEncoder( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - timbre_hidden_dim=timbre_hidden_dim, - num_timbre_encoder_hidden_layers=num_timbre_encoder_hidden_layers, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - head_dim=head_dim, - rope_theta=rope_theta, - attention_bias=attention_bias, - attention_dropout=attention_dropout, - rms_norm_eps=rms_norm_eps, - sliding_window=sliding_window, - ) - - # Learned null-condition embedding for classifier-free guidance, trained with - # `cfg_ratio=0.15` in the original model. Broadcast along the sequence dim when used. - self.null_condition_emb = nn.Parameter(torch.randn(1, 1, hidden_size)) - - # Silence latent — VAE-encoded audio-silence, stored as (1, T_long, timbre_hidden_dim). - # When no reference audio is provided, the pipeline slices `silence_latent[:, :timbre_fix_frame, :]` - # and feeds that to the timbre encoder. Passing literal zeros puts the timbre encoder - # OOD and produces drone-like audio (observed on all text2music outputs before this fix). - # The placeholder here is overwritten by the converter with the real encoded silence, - # so its shape just needs to match the timbre-encoder input: last dim is - # `timbre_hidden_dim` (so smaller test configs with `timbre_hidden_dim != 64` also load). - self.register_buffer( - "silence_latent", - torch.zeros(1, 15000, timbre_hidden_dim), - persistent=True, +# Deprecated aliases kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class AceStepLyricEncoder(_AceStepLyricEncoder): + def __new__(cls, *args, **kwargs): + deprecate( + "AceStepLyricEncoder", + "1.0.0", + "Importing `AceStepLyricEncoder` from `diffusers.pipelines.ace_step.modeling_ace_step` is deprecated. " + "Import it from `diffusers.models.condition_embedders.condition_embedder_ace_step` instead.", ) + return super().__new__(cls) - def forward( - self, - text_hidden_states: torch.FloatTensor, - text_attention_mask: torch.Tensor, - lyric_hidden_states: torch.FloatTensor, - lyric_attention_mask: torch.Tensor, - refer_audio_acoustic_hidden_states_packed: torch.FloatTensor, - refer_audio_order_mask: torch.LongTensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: - text_hidden_states = self.text_projector(text_hidden_states) - lyric_hidden_states = self.lyric_encoder( - inputs_embeds=lyric_hidden_states, attention_mask=lyric_attention_mask +class AceStepTimbreEncoder(_AceStepTimbreEncoder): + def __new__(cls, *args, **kwargs): + deprecate( + "AceStepTimbreEncoder", + "1.0.0", + "Importing `AceStepTimbreEncoder` from `diffusers.pipelines.ace_step.modeling_ace_step` is deprecated. " + "Import it from `diffusers.models.condition_embedders.condition_embedder_ace_step` instead.", ) + return super().__new__(cls) - timbre_embs_unpack, timbre_embs_mask = self.timbre_encoder( - refer_audio_acoustic_hidden_states_packed, refer_audio_order_mask - ) - encoder_hidden_states, encoder_attention_mask = _pack_sequences( - lyric_hidden_states, timbre_embs_unpack, lyric_attention_mask, timbre_embs_mask +class AceStepConditionEncoder(_AceStepConditionEncoder): + def __new__(cls, *args, **kwargs): + deprecate( + "AceStepConditionEncoder", + "1.0.0", + "Importing `AceStepConditionEncoder` from `diffusers.pipelines.ace_step.modeling_ace_step` is deprecated. " + "Import it with `from diffusers import AceStepConditionEncoder` instead.", ) - encoder_hidden_states, encoder_attention_mask = _pack_sequences( - encoder_hidden_states, text_hidden_states, encoder_attention_mask, text_attention_mask - ) - - return encoder_hidden_states, encoder_attention_mask + return super().__new__(cls) diff --git a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py index 8f8e2e1be95b..d566b6d5e305 100644 --- a/src/diffusers/pipelines/ace_step/pipeline_ace_step.py +++ b/src/diffusers/pipelines/ace_step/pipeline_ace_step.py @@ -22,12 +22,13 @@ from ...guiders.adaptive_projected_guidance import MomentumBuffer, normalized_guidance from ...loaders import AceStepLoraLoaderMixin from ...models import AutoencoderOobleck +from ...models.condition_embedders import AceStepConditionEncoder from ...models.transformers.ace_step_transformer import AceStepTransformer1DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .modeling_ace_step import AceStepAudioTokenDetokenizer, AceStepAudioTokenizer, AceStepConditionEncoder +from .modeling_ace_step import AceStepAudioTokenDetokenizer, AceStepAudioTokenizer logger = logging.get_logger(__name__) # pylint: disable=invalid-name diff --git a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py index dc2aa1df66a7..9a30dc5be62c 100644 --- a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass from typing import Any import torch @@ -28,138 +27,25 @@ AttnAddedKVProcessor, AttnProcessor, ) +from ...models.condition_embedders.condition_embedder_audioldm2 import ( + AudioLDM2ProjectionModel as _AudioLDM2ProjectionModel, +) +from ...models.condition_embedders.condition_embedder_audioldm2 import ( + AudioLDM2ProjectionModelOutput, # noqa: F401 + add_special_tokens, # noqa: F401 +) from ...models.embeddings import TimestepEmbedding, Timesteps from ...models.modeling_utils import ModelMixin from ...models.resnet import Downsample2D, ResnetBlock2D, Upsample2D from ...models.transformers.transformer_2d import Transformer2DModel from ...models.unets.unet_2d_blocks import DownBlock2D, UpBlock2D from ...models.unets.unet_2d_condition import UNet2DConditionOutput -from ...utils import BaseOutput, logging +from ...utils import deprecate, logging logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def add_special_tokens(hidden_states, attention_mask, sos_token, eos_token): - batch_size = hidden_states.shape[0] - - if attention_mask is not None: - # Add two more steps to attn mask - new_attn_mask_step = attention_mask.new_ones((batch_size, 1)) - attention_mask = torch.concat([new_attn_mask_step, attention_mask, new_attn_mask_step], dim=-1) - - # Add the SOS / EOS tokens at the start / end of the sequence respectively - sos_token = sos_token.expand(batch_size, 1, -1) - eos_token = eos_token.expand(batch_size, 1, -1) - hidden_states = torch.concat([sos_token, hidden_states, eos_token], dim=1) - return hidden_states, attention_mask - - -@dataclass -class AudioLDM2ProjectionModelOutput(BaseOutput): - """ - Args: - Class for AudioLDM2 projection layer's outputs. - hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): - Sequence of hidden-states obtained by linearly projecting the hidden-states for each of the text - encoders and subsequently concatenating them together. - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on padding token indices, formed by concatenating the attention masks - for the two text encoders together. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - """ - - hidden_states: torch.Tensor - attention_mask: torch.LongTensor | None = None - - -class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin): - """ - A simple linear projection model to map two text embeddings to a shared latent space. It also inserts learned - embedding vectors at the start and end of each text embedding sequence respectively. Each variable appended with - `_1` refers to that corresponding to the second text encoder. Otherwise, it is from the first. - - Args: - text_encoder_dim (`int`): - Dimensionality of the text embeddings from the first text encoder (CLAP). - text_encoder_1_dim (`int`): - Dimensionality of the text embeddings from the second text encoder (T5 or VITS). - langauge_model_dim (`int`): - Dimensionality of the text embeddings from the language model (GPT2). - """ - - @register_to_config - def __init__( - self, - text_encoder_dim, - text_encoder_1_dim, - langauge_model_dim, - use_learned_position_embedding=None, - max_seq_length=None, - ): - super().__init__() - # additional projection layers for each text encoder - self.projection = nn.Linear(text_encoder_dim, langauge_model_dim) - self.projection_1 = nn.Linear(text_encoder_1_dim, langauge_model_dim) - - # learnable SOS / EOS token embeddings for each text encoder - self.sos_embed = nn.Parameter(torch.ones(langauge_model_dim)) - self.eos_embed = nn.Parameter(torch.ones(langauge_model_dim)) - - self.sos_embed_1 = nn.Parameter(torch.ones(langauge_model_dim)) - self.eos_embed_1 = nn.Parameter(torch.ones(langauge_model_dim)) - - self.use_learned_position_embedding = use_learned_position_embedding - - # learable positional embedding for vits encoder - if self.use_learned_position_embedding is not None: - self.learnable_positional_embedding = torch.nn.Parameter( - torch.zeros((1, text_encoder_1_dim, max_seq_length)) - ) - - def forward( - self, - hidden_states: torch.Tensor | None = None, - hidden_states_1: torch.Tensor | None = None, - attention_mask: torch.LongTensor | None = None, - attention_mask_1: torch.LongTensor | None = None, - ): - hidden_states = self.projection(hidden_states) - hidden_states, attention_mask = add_special_tokens( - hidden_states, attention_mask, sos_token=self.sos_embed, eos_token=self.eos_embed - ) - - # Add positional embedding for Vits hidden state - if self.use_learned_position_embedding is not None: - hidden_states_1 = (hidden_states_1.permute(0, 2, 1) + self.learnable_positional_embedding).permute(0, 2, 1) - - hidden_states_1 = self.projection_1(hidden_states_1) - hidden_states_1, attention_mask_1 = add_special_tokens( - hidden_states_1, attention_mask_1, sos_token=self.sos_embed_1, eos_token=self.eos_embed_1 - ) - - # concatenate clap and t5 text encoding - hidden_states = torch.cat([hidden_states, hidden_states_1], dim=1) - - # concatenate attention masks - if attention_mask is None and attention_mask_1 is not None: - attention_mask = attention_mask_1.new_ones((hidden_states[:2])) - elif attention_mask is not None and attention_mask_1 is None: - attention_mask_1 = attention_mask.new_ones((hidden_states_1[:2])) - - if attention_mask is not None and attention_mask_1 is not None: - attention_mask = torch.cat([attention_mask, attention_mask_1], dim=-1) - else: - attention_mask = None - - return AudioLDM2ProjectionModelOutput( - hidden_states=hidden_states, - attention_mask=attention_mask, - ) - - class AudioLDM2UNet2DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin): r""" A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample @@ -1409,3 +1295,16 @@ def forward( hidden_states = upsampler(hidden_states, upsample_size) return hidden_states + + +# Deprecated alias kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class AudioLDM2ProjectionModel(_AudioLDM2ProjectionModel): + def __new__(cls, *args, **kwargs): + deprecate( + "AudioLDM2ProjectionModel", + "1.0.0", + "Importing `AudioLDM2ProjectionModel` from `diffusers.pipelines.audioldm2.modeling_audioldm2` is deprecated. " + "Import it with `from diffusers import AudioLDM2ProjectionModel` instead.", + ) + return super().__new__(cls) diff --git a/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py b/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py index b9e64c18c468..401cfe7fe0b3 100644 --- a/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py @@ -32,6 +32,7 @@ ) from ...models import AutoencoderKL +from ...models.condition_embedders import AudioLDM2ProjectionModel from ...schedulers import KarrasDiffusionSchedulers from ...utils import ( is_accelerate_available, @@ -43,7 +44,7 @@ from ...utils.import_utils import is_transformers_version from ...utils.torch_utils import empty_device_cache, randn_tensor from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .modeling_audioldm2 import AudioLDM2ProjectionModel, AudioLDM2UNet2DConditionModel +from .modeling_audioldm2 import AudioLDM2UNet2DConditionModel if is_librosa_available(): diff --git a/src/diffusers/pipelines/deprecated/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py b/src/diffusers/pipelines/deprecated/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py index d72d12a64945..00ebcda4defe 100644 --- a/src/diffusers/pipelines/deprecated/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +++ b/src/diffusers/pipelines/deprecated/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py @@ -30,6 +30,7 @@ from ....loaders import StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin from ....models import AutoencoderKL, UNet2DConditionModel from ....models.attention import GatedSelfAttentionDense +from ....models.condition_embedders import CLIPImageProjection from ....models.lora import adjust_lora_scale_text_encoder from ....schedulers import KarrasDiffusionSchedulers from ....utils import ( @@ -43,7 +44,6 @@ from ....utils.torch_utils import randn_tensor from ...pipeline_utils import DeprecatedPipelineMixin, DiffusionPipeline, StableDiffusionMixin from ...stable_diffusion import StableDiffusionPipelineOutput -from ...stable_diffusion.clip_image_project_model import CLIPImageProjection from ...stable_diffusion.safety_checker import StableDiffusionSafetyChecker diff --git a/src/diffusers/pipelines/flux/modeling_flux.py b/src/diffusers/pipelines/flux/modeling_flux.py index 916e3fbe5953..9e8647c3c0b3 100644 --- a/src/diffusers/pipelines/flux/modeling_flux.py +++ b/src/diffusers/pipelines/flux/modeling_flux.py @@ -12,35 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. - -from dataclasses import dataclass - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin -from ...utils import BaseOutput - - -@dataclass -class ReduxImageEncoderOutput(BaseOutput): - image_embeds: torch.Tensor | None = None - - -class ReduxImageEncoder(ModelMixin, ConfigMixin): - @register_to_config - def __init__( - self, - redux_dim: int = 1152, - txt_in_features: int = 4096, - ) -> None: - super().__init__() - - self.redux_up = nn.Linear(redux_dim, txt_in_features * 3) - self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features) - - def forward(self, x: torch.Tensor) -> ReduxImageEncoderOutput: - projected_x = self.redux_down(nn.functional.silu(self.redux_up(x))) - - return ReduxImageEncoderOutput(image_embeds=projected_x) +from ...models.condition_embedders.condition_embedder_flux_redux import ReduxImageEncoder as _ReduxImageEncoder +from ...models.condition_embedders.condition_embedder_flux_redux import ( + ReduxImageEncoderOutput, # noqa: F401 +) +from ...utils import deprecate + + +# Deprecated alias kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class ReduxImageEncoder(_ReduxImageEncoder): + def __new__(cls, *args, **kwargs): + deprecate( + "ReduxImageEncoder", + "1.0.0", + "Importing `ReduxImageEncoder` from `diffusers.pipelines.flux.modeling_flux` is deprecated. " + "Import it with `from diffusers import ReduxImageEncoder` instead.", + ) + return super().__new__(cls) diff --git a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py index e74d6d4e13b6..ee985b3165e8 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_prior_redux.py @@ -26,6 +26,7 @@ from ...image_processor import PipelineImageInput from ...loaders import FluxLoraLoaderMixin, TextualInversionLoaderMixin +from ...models.condition_embedders import ReduxImageEncoder from ...utils import ( USE_PEFT_BACKEND, is_torch_xla_available, @@ -36,7 +37,6 @@ ) from ...utils.torch_utils import get_module_execution_device from ..pipeline_utils import DiffusionPipeline -from .modeling_flux import ReduxImageEncoder from .pipeline_output import FluxPriorReduxPipelineOutput diff --git a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py index f579cf00dbe7..67e0593278c0 100644 --- a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py +++ b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py @@ -12,175 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -import torch - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin - - -class ResBlock(torch.nn.Module): - def __init__(self, channels: int, mid_channels: int | None = None, dims: int = 3): - super().__init__() - if mid_channels is None: - mid_channels = channels - - Conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d - - self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1) - self.norm1 = torch.nn.GroupNorm(32, mid_channels) - self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1) - self.norm2 = torch.nn.GroupNorm(32, channels) - self.activation = torch.nn.SiLU() - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - residual = hidden_states - hidden_states = self.conv1(hidden_states) - hidden_states = self.norm1(hidden_states) - hidden_states = self.activation(hidden_states) - hidden_states = self.conv2(hidden_states) - hidden_states = self.norm2(hidden_states) - hidden_states = self.activation(hidden_states + residual) - return hidden_states - - -class PixelShuffleND(torch.nn.Module): - def __init__(self, dims, upscale_factors=(2, 2, 2)): - super().__init__() - - self.dims = dims - self.upscale_factors = upscale_factors - - if dims not in [1, 2, 3]: - raise ValueError("dims must be 1, 2, or 3") - - def forward(self, x): - if self.dims == 3: - # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3) - return ( - x.unflatten(1, (-1, *self.upscale_factors[:3])) - .permute(0, 1, 5, 2, 6, 3, 7, 4) - .flatten(6, 7) - .flatten(4, 5) - .flatten(2, 3) - ) - elif self.dims == 2: - # spatial: b (c p1 p2) h w -> b c (h p1) (w p2) - return ( - x.unflatten(1, (-1, *self.upscale_factors[:2])).permute(0, 1, 4, 2, 5, 3).flatten(4, 5).flatten(2, 3) - ) - elif self.dims == 1: - # temporal: b (c p1) f h w -> b c (f p1) h w - return x.unflatten(1, (-1, *self.upscale_factors[:1])).permute(0, 1, 3, 2, 4, 5).flatten(2, 3) - - -class LTXLatentUpsamplerModel(ModelMixin, ConfigMixin): - """ - Model to spatially upsample VAE latents. - - Args: - in_channels (`int`, defaults to `128`): - Number of channels in the input latent - mid_channels (`int`, defaults to `512`): - Number of channels in the middle layers - num_blocks_per_stage (`int`, defaults to `4`): - Number of ResBlocks to use in each stage (pre/post upsampling) - dims (`int`, defaults to `3`): - Number of dimensions for convolutions (2 or 3) - spatial_upsample (`bool`, defaults to `True`): - Whether to spatially upsample the latent - temporal_upsample (`bool`, defaults to `False`): - Whether to temporally upsample the latent - """ - - @register_to_config - def __init__( - self, - in_channels: int = 128, - mid_channels: int = 512, - num_blocks_per_stage: int = 4, - dims: int = 3, - spatial_upsample: bool = True, - temporal_upsample: bool = False, - ): - super().__init__() - - self.in_channels = in_channels - self.mid_channels = mid_channels - self.num_blocks_per_stage = num_blocks_per_stage - self.dims = dims - self.spatial_upsample = spatial_upsample - self.temporal_upsample = temporal_upsample - - ConvNd = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d - - self.initial_conv = ConvNd(in_channels, mid_channels, kernel_size=3, padding=1) - self.initial_norm = torch.nn.GroupNorm(32, mid_channels) - self.initial_activation = torch.nn.SiLU() - - self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) - - if spatial_upsample and temporal_upsample: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(3), - ) - elif spatial_upsample: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(2), - ) - elif temporal_upsample: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(1), - ) - else: - raise ValueError("Either spatial_upsample or temporal_upsample must be True") - - self.post_upsample_res_blocks = torch.nn.ModuleList( - [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] +from ...models.latent_upscaler.latent_upscaler_ltx import LTXLatentUpsamplerModel as _LTXLatentUpsamplerModel +from ...models.latent_upscaler.latent_upscaler_ltx import ( + PixelShuffleND, # noqa: F401 + ResBlock, # noqa: F401 +) +from ...utils import deprecate + + +# Deprecated alias kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class LTXLatentUpsamplerModel(_LTXLatentUpsamplerModel): + def __new__(cls, *args, **kwargs): + deprecate( + "LTXLatentUpsamplerModel", + "1.0.0", + "Importing `LTXLatentUpsamplerModel` from `diffusers.pipelines.ltx.modeling_latent_upsampler` is deprecated. " + "Import it with `from diffusers import LTXLatentUpsamplerModel` instead.", ) - - self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - - if self.dims == 2: - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.initial_conv(hidden_states) - hidden_states = self.initial_norm(hidden_states) - hidden_states = self.initial_activation(hidden_states) - - for block in self.res_blocks: - hidden_states = block(hidden_states) - - hidden_states = self.upsampler(hidden_states) - - for block in self.post_upsample_res_blocks: - hidden_states = block(hidden_states) - - hidden_states = self.final_conv(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - else: - hidden_states = self.initial_conv(hidden_states) - hidden_states = self.initial_norm(hidden_states) - hidden_states = self.initial_activation(hidden_states) - - for block in self.res_blocks: - hidden_states = block(hidden_states) - - if self.temporal_upsample: - hidden_states = self.upsampler(hidden_states) - hidden_states = hidden_states[:, :, 1:, :, :] - else: - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.upsampler(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - - for block in self.post_upsample_res_blocks: - hidden_states = block(hidden_states) - - hidden_states = self.final_conv(hidden_states) - - return hidden_states + return super().__new__(cls) diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py index c6eff670908f..88da2ba99a9b 100644 --- a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py +++ b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py @@ -16,11 +16,11 @@ from ...image_processor import PipelineImageInput from ...models import AutoencoderKLLTXVideo +from ...models.latent_upscaler import LTXLatentUpsamplerModel from ...utils import get_logger from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .modeling_latent_upsampler import LTXLatentUpsamplerModel from .pipeline_output import LTXPipelineOutput diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py index 0b0b333209aa..56034b2866da 100644 --- a/src/diffusers/pipelines/ltx2/connectors.py +++ b/src/diffusers/pipelines/ltx2/connectors.py @@ -1,478 +1,36 @@ -import math - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin -from ...models.attention import FeedForward -from ...models.modeling_utils import ModelMixin -from ...models.transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoAttnProcessor - - -def per_layer_masked_mean_norm( - text_hidden_states: torch.Tensor, - sequence_lengths: torch.Tensor, - device: str | torch.device, - padding_side: str = "left", - scale_factor: int = 8, - eps: float = 1e-6, -): - """ - Performs per-batch per-layer normalization using a masked mean and range on per-layer text encoder hidden_states. - Respects the padding of the hidden states. - - Args: - text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`): - Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`). - sequence_lengths (`torch.Tensor of shape `(batch_size,)`): - The number of valid (non-padded) tokens for each batch instance. - device: (`str` or `torch.device`, *optional*): - torch device to place the resulting embeddings on - padding_side: (`str`, *optional*, defaults to `"left"`): - Whether the text tokenizer performs padding on the `"left"` or `"right"`. - scale_factor (`int`, *optional*, defaults to `8`): - Scaling factor to multiply the normalized hidden states by. - eps (`float`, *optional*, defaults to `1e-6`): - A small positive value for numerical stability when performing normalization. - - Returns: - `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`: - Normed and flattened text encoder hidden states. - """ - batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape - original_dtype = text_hidden_states.dtype - - # Create padding mask - token_indices = torch.arange(seq_len, device=device).unsqueeze(0) - if padding_side == "right": - # For right padding, valid tokens are from 0 to sequence_length-1 - mask = token_indices < sequence_lengths[:, None] # [batch_size, seq_len] - elif padding_side == "left": - # For left padding, valid tokens are from (T - sequence_length) to T-1 - start_indices = seq_len - sequence_lengths[:, None] # [batch_size, 1] - mask = token_indices >= start_indices # [B, T] - else: - raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") - mask = mask[:, :, None, None] # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1] - - # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len) - masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) - num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) - masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) - - # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len) - x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) - x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) - - # Normalization - normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) - normalized_hidden_states = normalized_hidden_states * scale_factor - - # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers) - normalized_hidden_states = normalized_hidden_states.flatten(2) - mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers) - normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) - normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype) - return normalized_hidden_states - - -def per_token_rms_norm(text_encoder_hidden_states: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: - variance = torch.mean(text_encoder_hidden_states**2, dim=2, keepdim=True) - norm_text_encoder_hidden_states = text_encoder_hidden_states * torch.rsqrt(variance + eps) - return norm_text_encoder_hidden_states - - -class LTX2RotaryPosEmbed1d(nn.Module): - """ - 1D rotary positional embeddings (RoPE) for the LTX 2.0 text encoder connectors. - """ - - def __init__( - self, - dim: int, - base_seq_len: int = 4096, - theta: float = 10000.0, - double_precision: bool = True, - rope_type: str = "interleaved", - num_attention_heads: int = 32, - ): - super().__init__() - if rope_type not in ["interleaved", "split"]: - raise ValueError(f"{rope_type=} not supported. Choose between 'interleaved' and 'split'.") - - self.dim = dim - self.base_seq_len = base_seq_len - self.theta = theta - self.double_precision = double_precision - self.rope_type = rope_type - self.num_attention_heads = num_attention_heads - - def forward( - self, - batch_size: int, - pos: int, - device: str | torch.device, - ) -> tuple[torch.Tensor, torch.Tensor]: - # 1. Get 1D position ids - grid_1d = torch.arange(pos, dtype=torch.float32, device=device) - # Get fractional indices relative to self.base_seq_len - grid_1d = grid_1d / self.base_seq_len - grid = grid_1d.unsqueeze(0).repeat(batch_size, 1) # [batch_size, seq_len] - - # 2. Calculate 1D RoPE frequencies - num_rope_elems = 2 # 1 (because 1D) * 2 (for cos, sin) = 2 - freqs_dtype = torch.float64 if self.double_precision else torch.float32 - pow_indices = torch.pow( - self.theta, - torch.linspace(start=0.0, end=1.0, steps=self.dim // num_rope_elems, dtype=freqs_dtype, device=device), +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ...models.condition_embedders.condition_embedder_ltx2 import ( + LTX2ConnectorTransformer1d, # noqa: F401 + LTX2RotaryPosEmbed1d, # noqa: F401 + LTX2TransformerBlock1d, # noqa: F401 + per_layer_masked_mean_norm, # noqa: F401 + per_token_rms_norm, # noqa: F401 +) +from ...models.condition_embedders.condition_embedder_ltx2 import LTX2TextConnectors as _LTX2TextConnectors +from ...utils import deprecate + + +# Deprecated alias kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class LTX2TextConnectors(_LTX2TextConnectors): + def __new__(cls, *args, **kwargs): + deprecate( + "LTX2TextConnectors", + "1.0.0", + "Importing `LTX2TextConnectors` from `diffusers.pipelines.ltx2.connectors` is deprecated. " + "Import it with `from diffusers import LTX2TextConnectors` instead.", ) - freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) - - # 3. Matrix-vector outer product between pos ids of shape (batch_size, seq_len) and freqs vector of shape - # (self.dim // 2,). - freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs # [B, seq_len, self.dim // 2] - - # 4. Get real, interleaved (cos, sin) frequencies, padded to self.dim - if self.rope_type == "interleaved": - cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) - sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) - - if self.dim % num_rope_elems != 0: - cos_padding = torch.ones_like(cos_freqs[:, :, : self.dim % num_rope_elems]) - sin_padding = torch.zeros_like(sin_freqs[:, :, : self.dim % num_rope_elems]) - cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) - sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) - - elif self.rope_type == "split": - expected_freqs = self.dim // 2 - current_freqs = freqs.shape[-1] - pad_size = expected_freqs - current_freqs - cos_freq = freqs.cos() - sin_freq = freqs.sin() - - if pad_size != 0: - cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) - sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) - - cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) - sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) - - # Reshape freqs to be compatible with multi-head attention - b = cos_freq.shape[0] - t = cos_freq.shape[1] - - cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1) - sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1) - - cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) - sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) - - return cos_freqs, sin_freqs - - -class LTX2TransformerBlock1d(nn.Module): - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - activation_fn: str = "gelu-approximate", - eps: float = 1e-6, - rope_type: str = "interleaved", - apply_gated_attention: bool = False, - ): - super().__init__() - - self.norm1 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) - self.attn1 = LTX2Attention( - query_dim=dim, - heads=num_attention_heads, - kv_heads=num_attention_heads, - dim_head=attention_head_dim, - rope_type=rope_type, - apply_gated_attention=apply_gated_attention, - processor=LTX2AudioVideoAttnProcessor(), - ) - - self.norm2 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) - self.ff = FeedForward(dim, activation_fn=activation_fn) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - rotary_emb: torch.Tensor | None = None, - ) -> torch.Tensor: - norm_hidden_states = self.norm1(hidden_states) - attn_hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, query_rotary_emb=rotary_emb) - hidden_states = hidden_states + attn_hidden_states - - norm_hidden_states = self.norm2(hidden_states) - ff_hidden_states = self.ff(norm_hidden_states) - hidden_states = hidden_states + ff_hidden_states - - return hidden_states - - -class LTX2ConnectorTransformer1d(nn.Module): - """ - A 1D sequence transformer for modalities such as text. - - In LTX 2.0, this is used to process the text encoder hidden states for each of the video and audio streams. - """ - - _supports_gradient_checkpointing = True - - def __init__( - self, - num_attention_heads: int = 30, - attention_head_dim: int = 128, - num_layers: int = 2, - num_learnable_registers: int | None = 128, - rope_base_seq_len: int = 4096, - rope_theta: float = 10000.0, - rope_double_precision: bool = True, - eps: float = 1e-6, - causal_temporal_positioning: bool = False, - rope_type: str = "interleaved", - gated_attention: bool = False, - ): - super().__init__() - self.num_attention_heads = num_attention_heads - self.inner_dim = num_attention_heads * attention_head_dim - self.causal_temporal_positioning = causal_temporal_positioning - - self.num_learnable_registers = num_learnable_registers - self.learnable_registers = None - if num_learnable_registers is not None: - init_registers = torch.rand(num_learnable_registers, self.inner_dim) * 2.0 - 1.0 - self.learnable_registers = torch.nn.Parameter(init_registers) - - self.rope = LTX2RotaryPosEmbed1d( - self.inner_dim, - base_seq_len=rope_base_seq_len, - theta=rope_theta, - double_precision=rope_double_precision, - rope_type=rope_type, - num_attention_heads=num_attention_heads, - ) - - self.transformer_blocks = torch.nn.ModuleList( - [ - LTX2TransformerBlock1d( - dim=self.inner_dim, - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - rope_type=rope_type, - apply_gated_attention=gated_attention, - ) - for _ in range(num_layers) - ] - ) - - self.norm_out = torch.nn.RMSNorm(self.inner_dim, eps=eps, elementwise_affine=False) - - self.gradient_checkpointing = False - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - attn_mask_binarize_threshold: float = -9000.0, - ) -> tuple[torch.Tensor, torch.Tensor]: - # hidden_states shape: [batch_size, seq_len, hidden_dim] - # attention_mask shape: [batch_size, seq_len] or [batch_size, 1, 1, seq_len] - batch_size, seq_len, _ = hidden_states.shape - - # 1. Replace padding with learned registers, if using - if self.learnable_registers is not None: - if seq_len % self.num_learnable_registers != 0: - raise ValueError( - f"The `hidden_states` sequence length {hidden_states.shape[1]} should be divisible by the number" - f" of learnable registers {self.num_learnable_registers}" - ) - - num_register_repeats = seq_len // self.num_learnable_registers - registers = ( - self.learnable_registers.unsqueeze(0).expand(num_register_repeats, -1, -1).reshape(seq_len, -1) - ) # [seq_len, inner_dim] - - binary_attn_mask = (attention_mask >= attn_mask_binarize_threshold).int() - if binary_attn_mask.ndim == 4: - binary_attn_mask = binary_attn_mask.squeeze(1).squeeze(1) # [B, 1, 1, L] --> [B, L] - - # Move the valid tokens to the front in their original order and fill the tail - # with registers indexed by absolute position, matching the original LTX - # implementation (`_replace_padded_with_learnable_registers`). A stable argsort - # of the inverted mask gathers valid tokens first while preserving their order. - order = torch.argsort(1 - binary_attn_mask, dim=1, stable=True) # [B, L] - front_aligned = torch.gather(hidden_states, 1, order.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1])) - num_valid = binary_attn_mask.sum(dim=1, keepdim=True) # [B, 1] - positions = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0) # [1, L] - front_mask = (positions < num_valid).unsqueeze(-1) # [B, L, 1] - registers_expanded = registers.unsqueeze(0).expand(batch_size, -1, -1) # [B, L, D] - hidden_states = torch.where(front_mask, front_aligned, registers_expanded.to(hidden_states.dtype)) - - # Overwrite attention_mask with an all-zeros mask if using registers. - attention_mask = torch.zeros_like(attention_mask) - - # 2. Calculate 1D RoPE positional embeddings - rotary_emb = self.rope(batch_size, seq_len, device=hidden_states.device) - - # 3. Run 1D transformer blocks - for block in self.transformer_blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(block, hidden_states, attention_mask, rotary_emb) - else: - hidden_states = block(hidden_states, attention_mask=attention_mask, rotary_emb=rotary_emb) - - hidden_states = self.norm_out(hidden_states) - - return hidden_states, attention_mask - - -class LTX2TextConnectors(ModelMixin, PeftAdapterMixin, ConfigMixin): - """ - Text connector stack used by LTX 2.0 to process the packed text encoder hidden states for both the video and audio - streams. - """ - - @register_to_config - def __init__( - self, - caption_channels: int = 3840, # default Gemma-3-12B text encoder hidden_size - text_proj_in_factor: int = 49, # num_layers + 1 for embedding layer = 48 + 1 for Gemma-3-12B - video_connector_num_attention_heads: int = 30, - video_connector_attention_head_dim: int = 128, - video_connector_num_layers: int = 2, - video_connector_num_learnable_registers: int | None = 128, - video_gated_attn: bool = False, - audio_connector_num_attention_heads: int = 30, - audio_connector_attention_head_dim: int = 128, - audio_connector_num_layers: int = 2, - audio_connector_num_learnable_registers: int | None = 128, - audio_gated_attn: bool = False, - connector_rope_base_seq_len: int = 4096, - rope_theta: float = 10000.0, - rope_double_precision: bool = True, - causal_temporal_positioning: bool = False, - rope_type: str = "interleaved", - per_modality_projections: bool = False, - video_hidden_dim: int = 4096, - audio_hidden_dim: int = 2048, - proj_bias: bool = False, - ): - super().__init__() - text_encoder_dim = caption_channels * text_proj_in_factor - if per_modality_projections: - self.video_text_proj_in = nn.Linear(text_encoder_dim, video_hidden_dim, bias=proj_bias) - self.audio_text_proj_in = nn.Linear(text_encoder_dim, audio_hidden_dim, bias=proj_bias) - else: - self.text_proj_in = nn.Linear(text_encoder_dim, caption_channels, bias=proj_bias) - - self.video_connector = LTX2ConnectorTransformer1d( - num_attention_heads=video_connector_num_attention_heads, - attention_head_dim=video_connector_attention_head_dim, - num_layers=video_connector_num_layers, - num_learnable_registers=video_connector_num_learnable_registers, - rope_base_seq_len=connector_rope_base_seq_len, - rope_theta=rope_theta, - rope_double_precision=rope_double_precision, - causal_temporal_positioning=causal_temporal_positioning, - rope_type=rope_type, - gated_attention=video_gated_attn, - ) - self.audio_connector = LTX2ConnectorTransformer1d( - num_attention_heads=audio_connector_num_attention_heads, - attention_head_dim=audio_connector_attention_head_dim, - num_layers=audio_connector_num_layers, - num_learnable_registers=audio_connector_num_learnable_registers, - rope_base_seq_len=connector_rope_base_seq_len, - rope_theta=rope_theta, - rope_double_precision=rope_double_precision, - causal_temporal_positioning=causal_temporal_positioning, - rope_type=rope_type, - gated_attention=audio_gated_attn, - ) - - def forward( - self, - text_encoder_hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - padding_side: str = "left", - scale_factor: int = 8, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Given per-layer text encoder hidden_states, extracts features and runs per-modality connectors to get text - embeddings for the LTX-2.X DiT models. - - Args: - text_encoder_hidden_states (`torch.Tensor`)): - Per-layer text encoder hidden_states. Can either be 4D with shape `(batch_size, seq_len, - caption_channels, text_proj_in_factor) or 3D with the last two dimensions flattened. - attention_mask (`torch.Tensor` of shape `(batch_size, seq_len)`): - Multiplicative binary attention mask where 1s indicate unmasked positions and 0s indicate masked - positions. - padding_side (`str`, *optional*, defaults to `"left"`): - The padding side used by the text encoder's text encoder (either `"left"` or `"right"`). Defaults to - `"left"` as this is what the default Gemma3-12B text encoder uses. Only used if - `per_modality_projections` is `False` (LTX-2.0 models). - scale_factor (`int`, *optional*, defaults to `8`): - Scale factor for masked mean/range normalization. Only used if `per_modality_projections` is `False` - (LTX-2.0 models). - """ - if text_encoder_hidden_states.ndim == 3: - # Ensure shape is [batch_size, seq_len, caption_channels, text_proj_in_factor] - text_encoder_hidden_states = text_encoder_hidden_states.unflatten(2, (self.config.caption_channels, -1)) - - if self.config.per_modality_projections: - # LTX-2.3 - norm_text_encoder_hidden_states = per_token_rms_norm(text_encoder_hidden_states) - - norm_text_encoder_hidden_states = norm_text_encoder_hidden_states.flatten(2, 3) - bool_mask = attention_mask.bool().unsqueeze(-1) - norm_text_encoder_hidden_states = torch.where( - bool_mask, norm_text_encoder_hidden_states, torch.zeros_like(norm_text_encoder_hidden_states) - ) - - # Rescale norms with respect to video and audio dims for feature extractors - video_scale_factor = math.sqrt(self.config.video_hidden_dim / self.config.caption_channels) - video_norm_text_emb = norm_text_encoder_hidden_states * video_scale_factor - audio_scale_factor = math.sqrt(self.config.audio_hidden_dim / self.config.caption_channels) - audio_norm_text_emb = norm_text_encoder_hidden_states * audio_scale_factor - - # Per-Modality Feature extractors - video_text_emb_proj = self.video_text_proj_in(video_norm_text_emb) - audio_text_emb_proj = self.audio_text_proj_in(audio_norm_text_emb) - else: - # LTX-2.0 - sequence_lengths = attention_mask.sum(dim=-1) - norm_text_encoder_hidden_states = per_layer_masked_mean_norm( - text_hidden_states=text_encoder_hidden_states, - sequence_lengths=sequence_lengths, - device=text_encoder_hidden_states.device, - padding_side=padding_side, - scale_factor=scale_factor, - ) - - text_emb_proj = self.text_proj_in(norm_text_encoder_hidden_states) - video_text_emb_proj = text_emb_proj - audio_text_emb_proj = text_emb_proj - - # Convert to additive attention mask for connectors - text_dtype = video_text_emb_proj.dtype - attention_mask = (attention_mask.to(torch.int64) - 1).to(text_dtype) - attention_mask = attention_mask.reshape(attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) - add_attn_mask = attention_mask * torch.finfo(text_dtype).max - - video_text_embedding, video_attn_mask = self.video_connector(video_text_emb_proj, add_attn_mask) - - # Convert video attn mask to binary (multiplicative) mask and mask video text embedding - binary_attn_mask = (video_attn_mask < 1e-6).to(torch.int64) - binary_attn_mask = binary_attn_mask.reshape(video_text_embedding.shape[0], video_text_embedding.shape[1], 1) - video_text_embedding = video_text_embedding * binary_attn_mask - - audio_text_embedding, _ = self.audio_connector(audio_text_emb_proj, add_attn_mask) - - return video_text_embedding, audio_text_embedding, binary_attn_mask.squeeze(-1) + return super().__new__(cls) diff --git a/src/diffusers/pipelines/ltx2/latent_upsampler.py b/src/diffusers/pipelines/ltx2/latent_upsampler.py index 329ced36d45b..faf8583c04ad 100644 --- a/src/diffusers/pipelines/ltx2/latent_upsampler.py +++ b/src/diffusers/pipelines/ltx2/latent_upsampler.py @@ -12,274 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -import math - -import torch -import torch.nn.functional as F - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin - - -RATIONAL_RESAMPLER_SCALE_MAPPING = { - 0.75: (3, 4), - 1.5: (3, 2), - 2.0: (2, 1), - 4.0: (4, 1), -} - - -# Copied from diffusers.pipelines.ltx.modeling_latent_upsampler.ResBlock -class ResBlock(torch.nn.Module): - def __init__(self, channels: int, mid_channels: int | None = None, dims: int = 3): - super().__init__() - if mid_channels is None: - mid_channels = channels - - Conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d - - self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1) - self.norm1 = torch.nn.GroupNorm(32, mid_channels) - self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1) - self.norm2 = torch.nn.GroupNorm(32, channels) - self.activation = torch.nn.SiLU() - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - residual = hidden_states - hidden_states = self.conv1(hidden_states) - hidden_states = self.norm1(hidden_states) - hidden_states = self.activation(hidden_states) - hidden_states = self.conv2(hidden_states) - hidden_states = self.norm2(hidden_states) - hidden_states = self.activation(hidden_states + residual) - return hidden_states - - -# Copied from diffusers.pipelines.ltx.modeling_latent_upsampler.PixelShuffleND -class PixelShuffleND(torch.nn.Module): - def __init__(self, dims, upscale_factors=(2, 2, 2)): - super().__init__() - - self.dims = dims - self.upscale_factors = upscale_factors - - if dims not in [1, 2, 3]: - raise ValueError("dims must be 1, 2, or 3") - - def forward(self, x): - if self.dims == 3: - # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3) - return ( - x.unflatten(1, (-1, *self.upscale_factors[:3])) - .permute(0, 1, 5, 2, 6, 3, 7, 4) - .flatten(6, 7) - .flatten(4, 5) - .flatten(2, 3) - ) - elif self.dims == 2: - # spatial: b (c p1 p2) h w -> b c (h p1) (w p2) - return ( - x.unflatten(1, (-1, *self.upscale_factors[:2])).permute(0, 1, 4, 2, 5, 3).flatten(4, 5).flatten(2, 3) - ) - elif self.dims == 1: - # temporal: b (c p1) f h w -> b c (f p1) h w - return x.unflatten(1, (-1, *self.upscale_factors[:1])).permute(0, 1, 3, 2, 4, 5).flatten(2, 3) - - -class BlurDownsample(torch.nn.Module): - """ - Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel. Applies only on H,W. - Works for dims=2 or dims=3 (per-frame). - """ - - def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None: - super().__init__() - - if dims not in (2, 3): - raise ValueError(f"`dims` must be either 2 or 3 but is {dims}") - if kernel_size < 3 or kernel_size % 2 != 1: - raise ValueError(f"`kernel_size` must be an odd number >= 3 but is {kernel_size}") - - self.dims = dims - self.stride = stride - self.kernel_size = kernel_size - - # 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from - # the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and - # provides a smooth approximation of a Gaussian filter (often called a "binomial filter"). - # The 2D kernel is constructed as the outer product and normalized. - k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)]) - k2d = k[:, None] @ k[None, :] - k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size) - self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.stride == 1: - return x - - if self.dims == 2: - c = x.shape[1] - weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise - x = F.conv2d(x, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c) - else: - # dims == 3: apply per-frame on H,W - b, c, f, _, _ = x.shape - x = x.transpose(1, 2).flatten(0, 1) # [B, C, F, H, W] --> [B * F, C, H, W] - - weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise - x = F.conv2d(x, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c) - - h2, w2 = x.shape[-2:] - x = x.unflatten(0, (b, f)).reshape(b, -1, f, h2, w2) # [B * F, C, H, W] --> [B, C, F, H, W] - return x - - -class SpatialRationalResampler(torch.nn.Module): - """ - Scales by the spatial size of the input by a rational number `scale`. For example, `scale = 0.75` will downsample - by a factor of 3 / 4, while `scale = 1.5` will upsample by a factor of 3 / 2. This works by first upsampling the - input by the (integer) numerator of `scale`, and then performing a blur + stride anti-aliased downsample by the - (integer) denominator. - """ - - def __init__(self, mid_channels: int = 1024, scale: float = 2.0): - super().__init__() - self.scale = float(scale) - num_denom = RATIONAL_RESAMPLER_SCALE_MAPPING.get(scale, None) - if num_denom is None: - raise ValueError( - f"The supplied `scale` {scale} is not supported; supported scales are {list(RATIONAL_RESAMPLER_SCALE_MAPPING.keys())}" - ) - self.num, self.den = num_denom - - self.conv = torch.nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1) - self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) - self.blur_down = BlurDownsample(dims=2, stride=self.den) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # Expected x shape: [B * F, C, H, W] - # b, _, f, h, w = x.shape - # x = x.transpose(1, 2).flatten(0, 1) # [B, C, F, H, W] --> [B * F, C, H, W] - x = self.conv(x) - x = self.pixel_shuffle(x) - x = self.blur_down(x) - # x = x.unflatten(0, (b, f)).reshape(b, -1, f, h, w) # [B * F, C, H, W] --> [B, C, F, H, W] - return x - - -class LTX2LatentUpsamplerModel(ModelMixin, ConfigMixin): - """ - Model to spatially upsample VAE latents. - - Args: - in_channels (`int`, defaults to `128`): - Number of channels in the input latent - mid_channels (`int`, defaults to `512`): - Number of channels in the middle layers - num_blocks_per_stage (`int`, defaults to `4`): - Number of ResBlocks to use in each stage (pre/post upsampling) - dims (`int`, defaults to `3`): - Number of dimensions for convolutions (2 or 3) - spatial_upsample (`bool`, defaults to `True`): - Whether to spatially upsample the latent - temporal_upsample (`bool`, defaults to `False`): - Whether to temporally upsample the latent - """ - - @register_to_config - def __init__( - self, - in_channels: int = 128, - mid_channels: int = 1024, - num_blocks_per_stage: int = 4, - dims: int = 3, - spatial_upsample: bool = True, - temporal_upsample: bool = False, - rational_spatial_scale: float = 2.0, - use_rational_resampler: bool = True, - ): - super().__init__() - - self.in_channels = in_channels - self.mid_channels = mid_channels - self.num_blocks_per_stage = num_blocks_per_stage - self.dims = dims - self.spatial_upsample = spatial_upsample - self.temporal_upsample = temporal_upsample - - ConvNd = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d - - self.initial_conv = ConvNd(in_channels, mid_channels, kernel_size=3, padding=1) - self.initial_norm = torch.nn.GroupNorm(32, mid_channels) - self.initial_activation = torch.nn.SiLU() - - self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) - - if spatial_upsample and temporal_upsample: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(3), - ) - elif spatial_upsample: - if use_rational_resampler: - self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=rational_spatial_scale) - else: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(2), - ) - elif temporal_upsample: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(1), - ) - else: - raise ValueError("Either spatial_upsample or temporal_upsample must be True") - - self.post_upsample_res_blocks = torch.nn.ModuleList( - [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] +from ...models.latent_upscaler.latent_upscaler_ltx2 import ( + RATIONAL_RESAMPLER_SCALE_MAPPING, # noqa: F401 + BlurDownsample, # noqa: F401 + PixelShuffleND, # noqa: F401 + ResBlock, # noqa: F401 + SpatialRationalResampler, # noqa: F401 +) +from ...models.latent_upscaler.latent_upscaler_ltx2 import LTX2LatentUpsamplerModel as _LTX2LatentUpsamplerModel +from ...utils import deprecate + + +# Deprecated alias kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class LTX2LatentUpsamplerModel(_LTX2LatentUpsamplerModel): + def __new__(cls, *args, **kwargs): + deprecate( + "LTX2LatentUpsamplerModel", + "1.0.0", + "Importing `LTX2LatentUpsamplerModel` from `diffusers.pipelines.ltx2.latent_upsampler` is deprecated. " + "Import it with `from diffusers import LTX2LatentUpsamplerModel` instead.", ) - - self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = hidden_states.shape - - if self.dims == 2: - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.initial_conv(hidden_states) - hidden_states = self.initial_norm(hidden_states) - hidden_states = self.initial_activation(hidden_states) - - for block in self.res_blocks: - hidden_states = block(hidden_states) - - hidden_states = self.upsampler(hidden_states) - - for block in self.post_upsample_res_blocks: - hidden_states = block(hidden_states) - - hidden_states = self.final_conv(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - else: - hidden_states = self.initial_conv(hidden_states) - hidden_states = self.initial_norm(hidden_states) - hidden_states = self.initial_activation(hidden_states) - - for block in self.res_blocks: - hidden_states = block(hidden_states) - - if self.temporal_upsample: - hidden_states = self.upsampler(hidden_states) - hidden_states = hidden_states[:, :, 1:, :, :] - else: - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.upsampler(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) - - for block in self.post_upsample_res_blocks: - hidden_states = block(hidden_states) - - hidden_states = self.final_conv(hidden_states) - - return hidden_states + return super().__new__(cls) diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py index baedfcd4b305..5dd3c3ca62e6 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2.py @@ -32,13 +32,13 @@ from ...image_processor import PipelineImageInput from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.condition_embedders import LTX2TextConnectors from ...models.transformers import LTX2VideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .connectors import LTX2TextConnectors from .duration_head import LTX2DurationHead from .pipeline_output import LTX2PipelineOutput from .prompt_enhancement import ( diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py index ef137ae0dc79..86c9ba8d4bd5 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_condition.py @@ -34,13 +34,13 @@ from ...image_processor import PipelineImageInput from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.condition_embedders import LTX2TextConnectors from ...models.transformers import LTX2VideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .connectors import LTX2TextConnectors from .duration_head import LTX2DurationHead from .pipeline_output import LTX2PipelineOutput from .prompt_enhancement import ( diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr.py index 936af16d8805..468872e8e253 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr.py @@ -32,13 +32,13 @@ from ...image_processor import PipelineImageInput from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.condition_embedders import LTX2TextConnectors from ...models.transformers import LTX2VideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler, LTXEulerAncestralRFScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .connectors import LTX2TextConnectors from .dfr_layout import LTX2DFREpilogueTile, resolve_canvas, video_tile_plan from .duration_head import LTX2DurationHead from .pipeline_ltx2_condition import LTX2VideoCondition diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr_temporal_refine.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr_temporal_refine.py index bdf9c8e67677..9309aa597309 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr_temporal_refine.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr_temporal_refine.py @@ -29,15 +29,15 @@ from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.condition_embedders import LTX2TextConnectors +from ...models.latent_upscaler import LTX2LatentUpsamplerModel from ...models.transformers import LTX2VideoTransformer3DModel from ...schedulers import LTXEulerAncestralRFScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .connectors import LTX2TextConnectors from .dfr_layout import temporal_tile_plan -from .latent_upsampler import LTX2LatentUpsamplerModel from .pipeline_ltx2_condition import LTX2VideoCondition from .pipeline_output import LTX2DFRPipelineOutput from .utils import ( @@ -101,7 +101,7 @@ def _audio_window_for_tile( ```py >>> import torch >>> from diffusers import LTX2DFRPipeline, LTX2DFRTemporalRefinePipeline, LTXEulerAncestralRFScheduler - >>> from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel + >>> from diffusers.models.latent_upscaler import LTX2LatentUpsamplerModel >>> pipe = LTX2DFRPipeline.from_pretrained("Lightricks/LTX-2.5-Diffusers", torch_dtype=torch.bfloat16) >>> temporal_upsampler = LTX2LatentUpsamplerModel.from_pretrained( diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_hdr_lora.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_hdr_lora.py index d9d95cdc0a0b..9f66e0355ff5 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_hdr_lora.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_hdr_lora.py @@ -30,12 +30,12 @@ from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.condition_embedders import LTX2TextConnectors from ...models.transformers import LTX2VideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline -from .connectors import LTX2TextConnectors from .image_processor import LTX2VideoHDRProcessor from .pipeline_output import LTX2PipelineOutput from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_ic_lora.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_ic_lora.py index 3d13ea4effa2..f405e22f7f2a 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_ic_lora.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_ic_lora.py @@ -34,13 +34,13 @@ from ...image_processor import PipelineImageInput from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.condition_embedders import LTX2TextConnectors from ...models.transformers import LTX2VideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .connectors import LTX2TextConnectors from .pipeline_ltx2_condition import LTX2VideoCondition from .pipeline_output import LTX2PipelineOutput from .prompt_enhancement import ( diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py index 7f88a56f16cf..e3c132df0365 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_image2video.py @@ -32,13 +32,13 @@ from ...image_processor import PipelineImageInput from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.condition_embedders import LTX2TextConnectors from ...models.transformers import LTX2VideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .connectors import LTX2TextConnectors from .duration_head import LTX2DurationHead from .pipeline_output import LTX2PipelineOutput from .prompt_enhancement import ( diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py index 09d842927d12..440d74529598 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_latent_upsample.py @@ -17,12 +17,12 @@ from ...image_processor import PipelineImageInput from ...models import AutoencoderKLLTX2Video +from ...models.latent_upscaler import LTX2LatentUpsamplerModel from ...utils import get_logger, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..ltx.pipeline_output import LTXPipelineOutput from ..pipeline_utils import DiffusionPipeline -from .latent_upsampler import LTX2LatentUpsamplerModel logger = get_logger(__name__) # pylint: disable=invalid-name @@ -34,7 +34,7 @@ >>> import torch >>> from diffusers import LTX2ImageToVideoPipeline, LTX2LatentUpsamplePipeline >>> from diffusers.utils import encode_video - >>> from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel + >>> from diffusers.models.latent_upscaler import LTX2LatentUpsamplerModel >>> from diffusers.utils import load_image >>> pipe = LTX2ImageToVideoPipeline.from_pretrained("Lightricks/LTX-2", torch_dtype=torch.bfloat16) diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 6958f49c8ddd..c781d5df0100 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -343,10 +343,30 @@ def _unwrap_model(model): return model +# Model classes that moved from a pipeline folder to `diffusers.models`. Existing `model_index.json` files still +# reference them as `[, ]`, so they are resolved to the top-level `diffusers` export +# instead of the deprecated alias left in the pipeline folder. +_RELOCATED_PIPELINE_CLASSES = { + ("ace_step", "AceStepConditionEncoder"), + ("audioldm2", "AudioLDM2ProjectionModel"), + ("flux", "ReduxImageEncoder"), + ("ltx", "LTXLatentUpsamplerModel"), + ("ltx2", "LTX2LatentUpsamplerModel"), + ("ltx2", "LTX2TextConnectors"), + ("stable_audio", "StableAudioProjectionModel"), + ("stable_audio_3", "StableAudio3DurationEmbedder"), + ("stable_diffusion", "CLIPImageProjection"), +} + + def maybe_raise_or_warn( library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module ): """Simple helper method to raise or warn in case incorrect module has been passed""" + if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: + library_name = "diffusers" + is_pipeline_module = False + if not is_pipeline_module: library = importlib.import_module(library_name) @@ -381,6 +401,9 @@ def maybe_raise_or_warn( def simple_get_class_obj(library_name, class_name): from diffusers import pipelines + if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: + library_name = "diffusers" + is_pipeline_module = hasattr(pipelines, library_name) if is_pipeline_module: @@ -414,6 +437,10 @@ def get_class_obj_and_candidates( if class_name.startswith("FlashPack"): class_name = class_name.removeprefix("FlashPack") + if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: + library_name = "diffusers" + is_pipeline_module = False + if is_pipeline_module: pipeline_module = getattr(pipelines, library_name) diff --git a/src/diffusers/pipelines/sana_wm/refiner.py b/src/diffusers/pipelines/sana_wm/refiner.py index b0dcc96c8efe..60134eb883ee 100644 --- a/src/diffusers/pipelines/sana_wm/refiner.py +++ b/src/diffusers/pipelines/sana_wm/refiner.py @@ -29,6 +29,7 @@ from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast from ...models.autoencoders import AutoencoderKLLTX2Video +from ...models.condition_embedders import LTX2TextConnectors from ...models.transformers.transformer_sana_wm_refiner import ( SanaWMLTX2RefinerTransformer3DModel, SanaWMRefinerKVCache, @@ -36,11 +37,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor - -# TODO: `LTX2TextConnectors` lives in the LTX-2 pipeline folder, so stage 2 has to reach across -# pipelines for it. Once https://github.com/huggingface/diffusers/issues/14749 moves the connector -# to a shared home (e.g. `models/`), import it from there and drop this cross-pipeline import. -from ..ltx2.connectors import LTX2TextConnectors from ..pipeline_utils import DiffusionPipeline diff --git a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py index d40269411bc0..12fa39096912 100644 --- a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py +++ b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py @@ -12,145 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass -from math import pi - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin -from ...utils import BaseOutput, logging - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class StableAudioPositionalEmbedding(nn.Module): - """Used for continuous time""" - - def __init__(self, dim: int): - super().__init__() - assert (dim % 2) == 0 - half_dim = dim // 2 - self.weights = nn.Parameter(torch.randn(half_dim)) - - def forward(self, times: torch.Tensor) -> torch.Tensor: - times = times[..., None] - freqs = times * self.weights[None] * 2 * pi - fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1) - fouriered = torch.cat((times, fouriered), dim=-1) - return fouriered - - -@dataclass -class StableAudioProjectionModelOutput(BaseOutput): - """ - Args: - Class for StableAudio projection layer's outputs. - text_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): - Sequence of hidden-states obtained by linearly projecting the hidden-states for the text encoder. - seconds_start_hidden_states (`torch.Tensor` of shape `(batch_size, 1, hidden_size)`, *optional*): - Sequence of hidden-states obtained by linearly projecting the audio start hidden states. - seconds_end_hidden_states (`torch.Tensor` of shape `(batch_size, 1, hidden_size)`, *optional*): - Sequence of hidden-states obtained by linearly projecting the audio end hidden states. - """ - - text_hidden_states: torch.Tensor | None = None - seconds_start_hidden_states: torch.Tensor | None = None - seconds_end_hidden_states: torch.Tensor | None = None - - -class StableAudioNumberConditioner(nn.Module): - """ - A simple linear projection model to map numbers to a latent space. - - Args: - number_embedding_dim (`int`): - Dimensionality of the number embeddings. - min_value (`int`): - The minimum value of the seconds number conditioning modules. - max_value (`int`): - The maximum value of the seconds number conditioning modules - internal_dim (`int`): - Dimensionality of the intermediate number hidden states. - """ - - def __init__( - self, - number_embedding_dim, - min_value, - max_value, - internal_dim: int | None = 256, - ): - super().__init__() - self.time_positional_embedding = nn.Sequential( - StableAudioPositionalEmbedding(internal_dim), - nn.Linear(in_features=internal_dim + 1, out_features=number_embedding_dim), - ) - - self.number_embedding_dim = number_embedding_dim - self.min_value = min_value - self.max_value = max_value - - def forward( - self, - floats: torch.Tensor, - ): - floats = floats.clamp(self.min_value, self.max_value) - - normalized_floats = (floats - self.min_value) / (self.max_value - self.min_value) - - # Cast floats to same type as embedder - embedder_dtype = next(self.time_positional_embedding.parameters()).dtype - normalized_floats = normalized_floats.to(embedder_dtype) - - embedding = self.time_positional_embedding(normalized_floats) - float_embeds = embedding.view(-1, 1, self.number_embedding_dim) - - return float_embeds - - -class StableAudioProjectionModel(ModelMixin, ConfigMixin): - """ - A simple linear projection model to map the conditioning values to a shared latent space. - - Args: - text_encoder_dim (`int`): - Dimensionality of the text embeddings from the text encoder (T5). - conditioning_dim (`int`): - Dimensionality of the output conditioning tensors. - min_value (`int`): - The minimum value of the seconds number conditioning modules. - max_value (`int`): - The maximum value of the seconds number conditioning modules - """ - - @register_to_config - def __init__(self, text_encoder_dim, conditioning_dim, min_value, max_value): - super().__init__() - self.text_projection = ( - nn.Identity() if conditioning_dim == text_encoder_dim else nn.Linear(text_encoder_dim, conditioning_dim) - ) - self.start_number_conditioner = StableAudioNumberConditioner(conditioning_dim, min_value, max_value) - self.end_number_conditioner = StableAudioNumberConditioner(conditioning_dim, min_value, max_value) - - def forward( - self, - text_hidden_states: torch.Tensor | None = None, - start_seconds: torch.Tensor | None = None, - end_seconds: torch.Tensor | None = None, - ): - text_hidden_states = ( - text_hidden_states if text_hidden_states is None else self.text_projection(text_hidden_states) - ) - seconds_start_hidden_states = ( - start_seconds if start_seconds is None else self.start_number_conditioner(start_seconds) - ) - seconds_end_hidden_states = end_seconds if end_seconds is None else self.end_number_conditioner(end_seconds) - - return StableAudioProjectionModelOutput( - text_hidden_states=text_hidden_states, - seconds_start_hidden_states=seconds_start_hidden_states, - seconds_end_hidden_states=seconds_end_hidden_states, +from ...models.condition_embedders.condition_embedder_stable_audio import ( + StableAudioNumberConditioner, # noqa: F401 + StableAudioPositionalEmbedding, # noqa: F401 + StableAudioProjectionModelOutput, # noqa: F401 +) +from ...models.condition_embedders.condition_embedder_stable_audio import ( + StableAudioProjectionModel as _StableAudioProjectionModel, +) +from ...utils import deprecate + + +# Deprecated alias kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class StableAudioProjectionModel(_StableAudioProjectionModel): + def __new__(cls, *args, **kwargs): + deprecate( + "StableAudioProjectionModel", + "1.0.0", + "Importing `StableAudioProjectionModel` from `diffusers.pipelines.stable_audio.modeling_stable_audio` is deprecated. " + "Import it with `from diffusers import StableAudioProjectionModel` instead.", ) + return super().__new__(cls) diff --git a/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py b/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py index b5f7cbfdf93f..0f5aa8658bf1 100644 --- a/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py +++ b/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py @@ -23,12 +23,12 @@ ) from ...models import AutoencoderOobleck, StableAudioDiTModel +from ...models.condition_embedders import StableAudioProjectionModel from ...models.embeddings import get_1d_rotary_pos_embed from ...schedulers import EDMDPMSolverMultistepScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .modeling_stable_audio import StableAudioProjectionModel if is_torch_xla_available(): diff --git a/src/diffusers/pipelines/stable_audio_3/modeling_stable_audio_3.py b/src/diffusers/pipelines/stable_audio_3/modeling_stable_audio_3.py index 52def1d8871d..105e3da10da2 100644 --- a/src/diffusers/pipelines/stable_audio_3/modeling_stable_audio_3.py +++ b/src/diffusers/pipelines/stable_audio_3/modeling_stable_audio_3.py @@ -12,98 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Auxiliary modules for the Stable Audio 3 pipeline. - -``StableAudio3DurationEmbedder`` mirrors the ``NumberConditioner`` used in the reference implementation of SA3 -(``stable_audio_tools``). It encodes a duration value (in seconds) into a fixed-dimensional vector that is passed to -the DiT as global conditioning (AdaLN). - -Architecture (matches reference ``conditioners.py`` ``NumberConditioner`` with ``fourier_features_type="expo"``): - - 1. Normalize: ``t_norm = clamp(t, min_val, max_val); t_norm = (t_norm - min_val) / (max_val - min_val)`` - 2. ``ExpoFourierFeatures``: fixed exponentially-spaced Fourier basis (ref: - ``stable_audio_tools.models.blocks.ExpoFourierFeatures``) - 3. Linear projection: ``fourier_dim → output_dim`` -""" - -import math - -import torch -import torch.nn as nn - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin -from ...utils import logging - - -logger = logging.get_logger(__name__) - - -class StableAudio3DurationEmbedder(ModelMixin, ConfigMixin): - """ - Embeds a duration value (in seconds) into a global conditioning vector for the Stable Audio 3 DiT (used as the - ``global_hidden_states`` AdaLN input). - - Replicates ``NumberConditioner(fourier_features_type="expo")`` from the SA3 reference implementation. - - Args: - output_dim (`int`, defaults to 768): - Dimension of the output embedding. Must match the DiT's ``global_cond_dim``. - fourier_dim (`int`, defaults to 256): - Internal Fourier feature dimension (must be even). - min_val (`float`, defaults to 0.0): - Minimum duration value for normalization clamping. - max_val (`float`, defaults to 384.0): - Maximum duration value for normalization clamping. Values above this are clamped. 384 seconds is the - production SA3 Medium upper bound for the ``seconds_total`` conditioner. - min_freq (`float`, defaults to 0.5): - Minimum frequency for the exponential Fourier basis. - max_freq (`float`, defaults to 10000.0): - Maximum frequency for the exponential Fourier basis. - """ - - @register_to_config - def __init__( - self, - output_dim: int = 768, - fourier_dim: int = 256, - min_val: float = 0.0, - max_val: float = 384.0, - min_freq: float = 0.5, - max_freq: float = 10000.0, - ) -> None: - super().__init__() - - # ExpoFourierFeatures — fixed (no learnable parameters) - # Frequencies are exponentially spaced from min_freq to max_freq. - half = fourier_dim // 2 - ramp = torch.linspace(0.0, 1.0, half) - log_min = math.log(min_freq) - log_max = math.log(max_freq) - freqs = torch.exp(ramp * (log_max - log_min) + log_min) - self.register_buffer("freqs", freqs) - - # Linear projection: fourier_dim → output_dim - self.linear = nn.Linear(fourier_dim, output_dim) - - def forward(self, seconds: torch.Tensor) -> torch.Tensor: - """ - Args: - seconds (`torch.Tensor` of shape `(batch,)`): - Duration values in seconds. - - Returns: - `torch.Tensor` of shape `(batch, output_dim)` — duration embeddings. - """ - # 1. Normalize to [0, 1] - seconds = seconds.float().clamp(self.config.min_val, self.config.max_val) - t_norm = (seconds - self.config.min_val) / (self.config.max_val - self.config.min_val) - - # 2. Exponential Fourier features — run in fp32 for stability - t_norm = t_norm.reshape(-1, 1) # (B, 1) - args = t_norm * self.freqs.unsqueeze(0) * 2.0 * math.pi # (B, half) - fourier = torch.cat([args.cos(), args.sin()], dim=-1) # (B, fourier_dim) - - # 3. Linear projection - return self.linear(fourier.to(self.dtype)) # (B, output_dim) +from ...models.condition_embedders.condition_embedder_stable_audio_3 import ( + StableAudio3DurationEmbedder as _StableAudio3DurationEmbedder, +) +from ...utils import deprecate + + +# Deprecated alias kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class StableAudio3DurationEmbedder(_StableAudio3DurationEmbedder): + def __new__(cls, *args, **kwargs): + deprecate( + "StableAudio3DurationEmbedder", + "1.0.0", + "Importing `StableAudio3DurationEmbedder` from `diffusers.pipelines.stable_audio_3.modeling_stable_audio_3` is deprecated. " + "Import it with `from diffusers import StableAudio3DurationEmbedder` instead.", + ) + return super().__new__(cls) diff --git a/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3.py b/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3.py index 11004b31b70b..0308fd9860d4 100644 --- a/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3.py +++ b/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3.py @@ -35,12 +35,12 @@ from transformers import GemmaTokenizer, GemmaTokenizerFast, T5GemmaEncoderModel from ...models.autoencoders.autoencoder_same import AutoencoderSAME +from ...models.condition_embedders import StableAudio3DurationEmbedder from ...models.transformers.transformer_stable_audio3 import StableAudio3DiTModel from ...schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .modeling_stable_audio_3 import StableAudio3DurationEmbedder if is_torch_xla_available(): diff --git a/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_audio2audio.py b/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_audio2audio.py index b84388adb12a..f6c1b0ad267b 100644 --- a/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_audio2audio.py +++ b/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_audio2audio.py @@ -29,12 +29,12 @@ from transformers import GemmaTokenizer, GemmaTokenizerFast, T5GemmaEncoderModel from ...models.autoencoders.autoencoder_same import AutoencoderSAME +from ...models.condition_embedders import StableAudio3DurationEmbedder from ...models.transformers.transformer_stable_audio3 import StableAudio3DiTModel from ...schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from ...utils import logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .modeling_stable_audio_3 import StableAudio3DurationEmbedder logger = logging.get_logger(__name__) diff --git a/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_inpaint.py b/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_inpaint.py index 14bafdadbab3..f691ce4a1f02 100644 --- a/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_inpaint.py +++ b/src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_inpaint.py @@ -32,12 +32,12 @@ from transformers import GemmaTokenizer, GemmaTokenizerFast, T5GemmaEncoderModel from ...models.autoencoders.autoencoder_same import AutoencoderSAME +from ...models.condition_embedders import StableAudio3DurationEmbedder from ...models.transformers.transformer_stable_audio3 import StableAudio3DiTModel from ...schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from ...utils import logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline -from .modeling_stable_audio_3 import StableAudio3DurationEmbedder logger = logging.get_logger(__name__) diff --git a/src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py b/src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py index 30dd90242d07..6bc1d2a3b2f7 100644 --- a/src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py +++ b/src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py @@ -12,18 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -from torch import nn +from ...models.condition_embedders.condition_embedder_clip_image import CLIPImageProjection as _CLIPImageProjection +from ...utils import deprecate -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin - -class CLIPImageProjection(ModelMixin, ConfigMixin): - @register_to_config - def __init__(self, hidden_size: int = 768): - super().__init__() - self.hidden_size = hidden_size - self.project = nn.Linear(self.hidden_size, self.hidden_size, bias=False) - - def forward(self, x): - return self.project(x) +# Deprecated alias kept for backward compatibility. The warning is raised from `__new__` rather than `__init__` so +# that the parent's `__init__` signature, which `ConfigMixin` inspects when loading a config, stays intact. +class CLIPImageProjection(_CLIPImageProjection): + def __new__(cls, *args, **kwargs): + deprecate( + "CLIPImageProjection", + "1.0.0", + "Importing `CLIPImageProjection` from `diffusers.pipelines.stable_diffusion.clip_image_project_model` is deprecated. " + "Import it with `from diffusers import CLIPImageProjection` instead.", + ) + return super().__new__(cls) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 22e653a403cc..b13d46f68ff5 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -439,6 +439,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AceStepConditionEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AceStepTransformer1DModel(metaclass=DummyObject): _backends = ["torch"] @@ -544,6 +559,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AudioLDM2ProjectionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AuraFlowTransformer2DModel(metaclass=DummyObject): _backends = ["torch"] @@ -1099,6 +1129,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class CLIPImageProjection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class CogVideoXTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -1819,6 +1864,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class LTX2LatentUpsamplerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class LTX2TextConnectors(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class LTX2VideoDiffusionDecoderModel(metaclass=DummyObject): _backends = ["torch"] @@ -1849,6 +1924,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class LTXLatentUpsamplerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class LTXVideoTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -2254,6 +2344,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ReduxImageEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class RFD4ProteinaAutoencoder(metaclass=DummyObject): _backends = ["torch"] @@ -2449,6 +2554,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class StableAudio3DurationEmbedder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class StableAudioDiTModel(metaclass=DummyObject): _backends = ["torch"] @@ -2464,6 +2584,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class StableAudioProjectionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class T2IAdapter(metaclass=DummyObject): _backends = ["torch"] @@ -3066,21 +3201,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class CLIPImageProjection(metaclass=DummyObject): - _backends = ["torch"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch"]) - - class ConsistencyModelPipeline(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index c9bd3b496d60..4b6f0876b06d 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -1232,21 +1232,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class AceStepConditionEncoder(metaclass=DummyObject): - _backends = ["torch", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - class AceStepPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -1502,21 +1487,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class AudioLDM2ProjectionModel(metaclass=DummyObject): - _backends = ["torch", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - class AudioLDM2UNet2DConditionModel(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -1682,21 +1652,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class CLIPImageProjection(metaclass=DummyObject): - _backends = ["torch", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - class CogVideoXFunControlPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -4052,21 +4007,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class ReduxImageEncoder(metaclass=DummyObject): - _backends = ["torch", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - class SanaControlNetPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -4352,21 +4292,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class StableAudio3DurationEmbedder(metaclass=DummyObject): - _backends = ["torch", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - class StableAudio3InpaintPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -4412,21 +4337,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class StableAudioProjectionModel(metaclass=DummyObject): - _backends = ["torch", "transformers"] - - def __init__(self, *args, **kwargs): - requires_backends(self, ["torch", "transformers"]) - - @classmethod - def from_config(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - @classmethod - def from_pretrained(cls, *args, **kwargs): - requires_backends(cls, ["torch", "transformers"]) - - class StableCascadeCombinedPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/test_relocated_class_loading.py b/tests/models/test_relocated_class_loading.py new file mode 100644 index 000000000000..f232ce813f9d --- /dev/null +++ b/tests/models/test_relocated_class_loading.py @@ -0,0 +1,85 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import warnings + +import pytest +import torch + +import diffusers +from diffusers import pipelines +from diffusers.pipelines.pipeline_loading_utils import ( + ALL_IMPORTABLE_CLASSES, + get_class_obj_and_candidates, + simple_get_class_obj, +) + + +# (pipeline folder, deprecated module path, class name, init kwargs) +RELOCATED_CLASSES = [ + ("ace_step", "diffusers.pipelines.ace_step.modeling_ace_step", "AceStepConditionEncoder", {}), + ( + "audioldm2", + "diffusers.pipelines.audioldm2.modeling_audioldm2", + "AudioLDM2ProjectionModel", + {"text_encoder_dim": 8, "text_encoder_1_dim": 8, "langauge_model_dim": 16}, + ), + ("flux", "diffusers.pipelines.flux.modeling_flux", "ReduxImageEncoder", {}), + ("ltx", "diffusers.pipelines.ltx.modeling_latent_upsampler", "LTXLatentUpsamplerModel", {}), + ("ltx2", "diffusers.pipelines.ltx2.latent_upsampler", "LTX2LatentUpsamplerModel", {}), + ("ltx2", "diffusers.pipelines.ltx2.connectors", "LTX2TextConnectors", {}), + ( + "stable_audio", + "diffusers.pipelines.stable_audio.modeling_stable_audio", + "StableAudioProjectionModel", + {"text_encoder_dim": 8, "conditioning_dim": 8, "min_value": 0, "max_value": 10}, + ), + ( + "stable_audio_3", + "diffusers.pipelines.stable_audio_3.modeling_stable_audio_3", + "StableAudio3DurationEmbedder", + {}, + ), + ("stable_diffusion", "diffusers.pipelines.stable_diffusion.clip_image_project_model", "CLIPImageProjection", {}), +] +IDS = [name for _, _, name, _ in RELOCATED_CLASSES] + + +def _future_warnings(fn): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fn() + return [str(w.message) for w in caught if issubclass(w.category, FutureWarning)] + + +@pytest.mark.parametrize("folder, module, name, kwargs", RELOCATED_CLASSES, ids=IDS) +def test_deprecated_path_warns_and_subclasses_new_class(folder, module, name, kwargs): + new_cls = getattr(diffusers, name) + old_cls = getattr(importlib.import_module(module), name) + assert new_cls.__module__.startswith("diffusers.models.") + assert issubclass(old_cls, new_cls) + + with torch.device("meta"): + assert not _future_warnings(lambda: new_cls(**kwargs)) + messages = _future_warnings(lambda: old_cls(**kwargs)) + assert any("deprecated" in message for message in messages), messages + + +@pytest.mark.parametrize("folder, module, name, kwargs", RELOCATED_CLASSES, ids=IDS) +def test_model_index_entry_resolves_to_new_class(folder, module, name, kwargs): + new_cls = getattr(diffusers, name) + class_obj, _ = get_class_obj_and_candidates(folder, name, ALL_IMPORTABLE_CLASSES, pipelines, True) + assert class_obj is new_cls + assert simple_get_class_obj(folder, name) is new_cls diff --git a/tests/pipelines/ace_step/test_ace_step.py b/tests/pipelines/ace_step/test_ace_step.py index 5ced68691272..f48042275f04 100644 --- a/tests/pipelines/ace_step/test_ace_step.py +++ b/tests/pipelines/ace_step/test_ace_step.py @@ -20,12 +20,11 @@ import torch from transformers import AutoTokenizer, Qwen3Config, Qwen3Model -from diffusers import AutoencoderOobleck, FlowMatchEulerDiscreteScheduler +from diffusers import AceStepConditionEncoder, AutoencoderOobleck, FlowMatchEulerDiscreteScheduler from diffusers.models.transformers.ace_step_transformer import AceStepTransformer1DModel from diffusers.pipelines.ace_step import ( AceStepAudioTokenDetokenizer, AceStepAudioTokenizer, - AceStepConditionEncoder, AceStepPipeline, ) diff --git a/tests/pipelines/ltx/test_ltx_latent_upsample.py b/tests/pipelines/ltx/test_ltx_latent_upsample.py index 14622b59ee5d..65937fa74bd3 100644 --- a/tests/pipelines/ltx/test_ltx_latent_upsample.py +++ b/tests/pipelines/ltx/test_ltx_latent_upsample.py @@ -16,7 +16,7 @@ import torch from diffusers import AutoencoderKLLTXVideo, LTXLatentUpsamplePipeline -from diffusers.pipelines.ltx.modeling_latent_upsampler import LTXLatentUpsamplerModel +from diffusers.models.latent_upscaler import LTXLatentUpsamplerModel from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin diff --git a/tests/pipelines/ltx2/test_ltx2_connectors.py b/tests/pipelines/ltx2/test_ltx2_connectors.py index 702c20d86ee9..34972838c451 100644 --- a/tests/pipelines/ltx2/test_ltx2_connectors.py +++ b/tests/pipelines/ltx2/test_ltx2_connectors.py @@ -14,7 +14,7 @@ import torch -from diffusers.pipelines.ltx2.connectors import LTX2ConnectorTransformer1d +from diffusers.models.condition_embedders.condition_embedder_ltx2 import LTX2ConnectorTransformer1d from ...testing_utils import assert_tensors_close, enable_full_determinism diff --git a/tests/pipelines/ltx2/testing_utils.py b/tests/pipelines/ltx2/testing_utils.py index e3b99f56a5e7..1e3a2e225ebf 100644 --- a/tests/pipelines/ltx2/testing_utils.py +++ b/tests/pipelines/ltx2/testing_utils.py @@ -31,7 +31,9 @@ LTX2VideoTransformer3DModel, LTXEulerAncestralRFScheduler, ) -from diffusers.pipelines.ltx2 import LTX2DurationHead, LTX2LatentUpsamplerModel, LTX2TextConnectors +from diffusers.models.condition_embedders import LTX2TextConnectors +from diffusers.models.latent_upscaler import LTX2LatentUpsamplerModel +from diffusers.pipelines.ltx2 import LTX2DurationHead from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from ...testing_utils import torch_device diff --git a/tests/pipelines/stable_audio_3/test_stable_audio_3.py b/tests/pipelines/stable_audio_3/test_stable_audio_3.py index e38fed14738e..04daff7c1589 100644 --- a/tests/pipelines/stable_audio_3/test_stable_audio_3.py +++ b/tests/pipelines/stable_audio_3/test_stable_audio_3.py @@ -25,7 +25,7 @@ StableAudio3DiTModel, StableAudio3Pipeline, ) -from diffusers.pipelines.stable_audio_3.modeling_stable_audio_3 import StableAudio3DurationEmbedder +from diffusers.models.condition_embedders import StableAudio3DurationEmbedder from ...testing_utils import ( enable_full_determinism, From ee1c09a3a6de138eecd6fac6a870d547ec8d1cb4 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 23 Oct 2026 06:05:03 +0000 Subject: [PATCH 2/5] fix docstring inconsistencies. --- .../condition_embedder_ace_step.py | 42 +++++++++++++++++++ .../condition_embedder_audioldm2.py | 17 +++++++- .../condition_embedder_clip_image.py | 11 ++++- .../condition_embedder_flux_redux.py | 9 ++++ .../condition_embedder_ltx2.py | 4 ++ .../condition_embedder_stable_audio.py | 17 +++++++- .../latent_upscaler/latent_upscaler_ltx.py | 10 ++++- .../latent_upscaler/latent_upscaler_ltx2.py | 10 ++++- 8 files changed, 114 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py b/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py index d918cc6d02ec..8c753850a656 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py @@ -183,6 +183,16 @@ def forward( inputs_embeds: torch.FloatTensor, attention_mask: torch.Tensor, ) -> torch.Tensor: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, text_hidden_dim)`): + Lyric embeddings from the text encoder. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`): + Mask with 1 for valid tokens and 0 for padding. + + Returns: + `torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`: The encoded lyric hidden states. + """ inputs_embeds = self.embed_tokens(inputs_embeds) seq_len = inputs_embeds.shape[1] @@ -322,6 +332,18 @@ def forward( refer_audio_acoustic_hidden_states_packed: torch.FloatTensor, refer_audio_order_mask: torch.LongTensor, ) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + refer_audio_acoustic_hidden_states_packed (`torch.FloatTensor` of shape `(num_references, sequence_length, timbre_hidden_dim)`): + VAE latents of the reference audio clips, packed across the batch. + refer_audio_order_mask (`torch.LongTensor` of shape `(num_references,)`): + Index of the batch item each reference clip belongs to. + + Returns: + `Tuple[torch.Tensor, torch.Tensor]`: The timbre embeddings of shape `(batch_size, max_references, + hidden_size)` and a mask of shape `(batch_size, max_references)` with 1 for valid embeddings and 0 for + padding. + """ inputs_embeds = self.embed_tokens(refer_audio_acoustic_hidden_states_packed) seq_len = inputs_embeds.shape[1] @@ -457,6 +479,26 @@ def forward( refer_audio_acoustic_hidden_states_packed: torch.FloatTensor, refer_audio_order_mask: torch.LongTensor, ) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + text_hidden_states (`torch.FloatTensor` of shape `(batch_size, text_length, text_hidden_dim)`): + Prompt embeddings from the text encoder. + text_attention_mask (`torch.Tensor` of shape `(batch_size, text_length)`): + Mask with 1 for valid prompt tokens and 0 for padding. + lyric_hidden_states (`torch.FloatTensor` of shape `(batch_size, lyric_length, text_hidden_dim)`): + Lyric embeddings from the text encoder. + lyric_attention_mask (`torch.Tensor` of shape `(batch_size, lyric_length)`): + Mask with 1 for valid lyric tokens and 0 for padding. + refer_audio_acoustic_hidden_states_packed (`torch.FloatTensor` of shape `(num_references, sequence_length, timbre_hidden_dim)`): + VAE latents of the reference audio clips, packed across the batch. + refer_audio_order_mask (`torch.LongTensor` of shape `(num_references,)`): + Index of the batch item each reference clip belongs to. + + Returns: + `Tuple[torch.Tensor, torch.Tensor]`: The packed lyric, timbre, and prompt hidden states of shape + `(batch_size, sequence_length, hidden_size)`, with valid tokens first, and their attention mask of shape + `(batch_size, sequence_length)`. + """ text_hidden_states = self.text_projector(text_hidden_states) lyric_hidden_states = self.lyric_encoder( diff --git a/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py b/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py index 4fb7f62867ba..616b885ba822 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py @@ -107,7 +107,22 @@ def forward( hidden_states_1: torch.Tensor | None = None, attention_mask: torch.LongTensor | None = None, attention_mask_1: torch.LongTensor | None = None, - ): + ) -> AudioLDM2ProjectionModelOutput: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, text_encoder_dim)`, *optional*): + Text embeddings from the first text encoder (CLAP). + hidden_states_1 (`torch.Tensor` of shape `(batch_size, sequence_length_1, text_encoder_1_dim)`, *optional*): + Text embeddings from the second text encoder (T5 or VITS). + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Attention mask for `hidden_states`. + attention_mask_1 (`torch.LongTensor` of shape `(batch_size, sequence_length_1)`, *optional*): + Attention mask for `hidden_states_1`. + + Returns: + [`AudioLDM2ProjectionModelOutput`]: The projected text embeddings of both text encoders, concatenated along + the sequence dimension, and their combined attention mask. + """ hidden_states = self.projection(hidden_states) hidden_states, attention_mask = add_special_tokens( hidden_states, attention_mask, sos_token=self.sos_embed, eos_token=self.eos_embed diff --git a/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py b/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py index a373f69fa39f..f717ff5f51d5 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import torch from torch import nn from ...configuration_utils import ConfigMixin, register_to_config @@ -25,5 +26,13 @@ def __init__(self, hidden_size: int = 768): self.hidden_size = hidden_size self.project = nn.Linear(self.hidden_size, self.hidden_size, bias=False) - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + r""" + Args: + x (`torch.Tensor` of shape `(..., hidden_size)`): + CLIP image embeddings to project. + + Returns: + `torch.Tensor` of shape `(..., hidden_size)`: The projected embeddings. + """ return self.project(x) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py b/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py index 526a8ddba99c..e8bae87fdafa 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py @@ -41,6 +41,15 @@ def __init__( self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features) def forward(self, x: torch.Tensor) -> ReduxImageEncoderOutput: + r""" + Args: + x (`torch.Tensor` of shape `(batch_size, sequence_length, redux_dim)`): + Image features from the image encoder. + + Returns: + [`ReduxImageEncoderOutput`]: The image embeddings of shape `(batch_size, sequence_length, + txt_in_features)`. + """ projected_x = self.redux_down(nn.functional.silu(self.redux_up(x))) return ReduxImageEncoderOutput(image_embeds=projected_x) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py b/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py index c119b6a0f6d0..3e5f2d269b08 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py @@ -421,6 +421,10 @@ def forward( scale_factor (`int`, *optional*, defaults to `8`): Scale factor for masked mean/range normalization. Only used if `per_modality_projections` is `False` (LTX-2.0 models). + + Returns: + `tuple[torch.Tensor, torch.Tensor, torch.Tensor]`: The video text embeddings, the audio text embeddings, + and a binary attention mask of shape `(batch_size, seq_len)` for the video text embeddings. """ if text_encoder_hidden_states.ndim == 3: # Ensure shape is [batch_size, seq_len, caption_channels, text_proj_in_factor] diff --git a/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py b/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py index 28bf56315939..0dac6ca01630 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py @@ -96,7 +96,7 @@ def __init__( def forward( self, floats: torch.Tensor, - ): + ) -> torch.Tensor: floats = floats.clamp(self.min_value, self.max_value) normalized_floats = (floats - self.min_value) / (self.max_value - self.min_value) @@ -140,7 +140,20 @@ def forward( text_hidden_states: torch.Tensor | None = None, start_seconds: torch.Tensor | None = None, end_seconds: torch.Tensor | None = None, - ): + ) -> StableAudioProjectionModelOutput: + r""" + Args: + text_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, text_encoder_dim)`, *optional*): + Text embeddings from the text encoder. + start_seconds (`torch.Tensor` of shape `(batch_size,)`, *optional*): + Audio start times in seconds. + end_seconds (`torch.Tensor` of shape `(batch_size,)`, *optional*): + Audio end times in seconds. + + Returns: + [`StableAudioProjectionModelOutput`]: The projected text, start-time, and end-time hidden states. Inputs + passed as `None` are returned as `None`. + """ text_hidden_states = ( text_hidden_states if text_hidden_states is None else self.text_projection(text_hidden_states) ) diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py index 76bb35ead842..f93ddddf13ec 100644 --- a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py +++ b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py @@ -53,7 +53,7 @@ def __init__(self, dims, upscale_factors=(2, 2, 2)): if dims not in [1, 2, 3]: raise ValueError("dims must be 1, 2, or 3") - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: if self.dims == 3: # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3) return ( @@ -144,6 +144,14 @@ def __init__( self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, in_channels, num_frames, height, width)`): + VAE latents to upsample. + + Returns: + `torch.Tensor`: The upsampled latents with `in_channels` channels. + """ batch_size, num_channels, num_frames, height, width = hidden_states.shape if self.dims == 2: diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py index d43a77920166..da3087a5b7ad 100644 --- a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py +++ b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py @@ -66,7 +66,7 @@ def __init__(self, dims, upscale_factors=(2, 2, 2)): if dims not in [1, 2, 3]: raise ValueError("dims must be 1, 2, or 3") - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: if self.dims == 3: # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3) return ( @@ -243,6 +243,14 @@ def __init__( self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, in_channels, num_frames, height, width)`): + VAE latents to upsample. + + Returns: + `torch.Tensor`: The upsampled latents with `in_channels` channels. + """ batch_size, num_channels, num_frames, height, width = hidden_states.shape if self.dims == 2: From e32c30972e35c5d123cd1b7e0b02442bf4302e7d Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 24 Oct 2026 03:28:42 +0000 Subject: [PATCH 3/5] rename to condition_embedder_gligen. --- src/diffusers/models/__init__.py | 2 +- src/diffusers/models/condition_embedders/__init__.py | 2 +- ...tion_embedder_clip_image.py => condition_embedder_gligen.py} | 0 .../pipelines/stable_diffusion/clip_image_project_model.py | 2 +- 4 files changed, 3 insertions(+), 3 deletions(-) rename src/diffusers/models/condition_embedders/{condition_embedder_clip_image.py => condition_embedder_gligen.py} (100%) diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index c29c178b83c5..4fa43aeee413 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -72,8 +72,8 @@ _import_structure["condition_embedders.condition_embedder_ace_step"] = ["AceStepConditionEncoder"] _import_structure["condition_embedders.condition_embedder_anima"] = ["AnimaTextConditioner"] _import_structure["condition_embedders.condition_embedder_audioldm2"] = ["AudioLDM2ProjectionModel"] - _import_structure["condition_embedders.condition_embedder_clip_image"] = ["CLIPImageProjection"] _import_structure["condition_embedders.condition_embedder_flux_redux"] = ["ReduxImageEncoder"] + _import_structure["condition_embedders.condition_embedder_gligen"] = ["CLIPImageProjection"] _import_structure["condition_embedders.condition_embedder_ltx2"] = ["LTX2TextConnectors"] _import_structure["condition_embedders.condition_embedder_minimax_music3"] = ["MiniMaxMusic3ConditionEncoder"] _import_structure["condition_embedders.condition_embedder_stable_audio"] = ["StableAudioProjectionModel"] diff --git a/src/diffusers/models/condition_embedders/__init__.py b/src/diffusers/models/condition_embedders/__init__.py index 7ba87e218095..7cc6519309de 100644 --- a/src/diffusers/models/condition_embedders/__init__.py +++ b/src/diffusers/models/condition_embedders/__init__.py @@ -5,8 +5,8 @@ from .condition_embedder_ace_step import AceStepConditionEncoder from .condition_embedder_anima import AnimaTextConditioner from .condition_embedder_audioldm2 import AudioLDM2ProjectionModel - from .condition_embedder_clip_image import CLIPImageProjection from .condition_embedder_flux_redux import ReduxImageEncoder + from .condition_embedder_gligen import CLIPImageProjection from .condition_embedder_ltx2 import LTX2TextConnectors from .condition_embedder_minimax_music3 import MiniMaxMusic3ConditionEncoder from .condition_embedder_stable_audio import StableAudioProjectionModel diff --git a/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py b/src/diffusers/models/condition_embedders/condition_embedder_gligen.py similarity index 100% rename from src/diffusers/models/condition_embedders/condition_embedder_clip_image.py rename to src/diffusers/models/condition_embedders/condition_embedder_gligen.py diff --git a/src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py b/src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py index 6bc1d2a3b2f7..d6473baaacb9 100644 --- a/src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py +++ b/src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ...models.condition_embedders.condition_embedder_clip_image import CLIPImageProjection as _CLIPImageProjection +from ...models.condition_embedders.condition_embedder_gligen import CLIPImageProjection as _CLIPImageProjection from ...utils import deprecate From 4042d45c241c5fe416c1b12edc5063c9c53b91fa Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 24 Oct 2026 03:53:23 +0000 Subject: [PATCH 4/5] remove exports. --- src/diffusers/pipelines/ace_step/modeling_ace_step.py | 5 +---- src/diffusers/pipelines/audioldm2/modeling_audioldm2.py | 4 ---- src/diffusers/pipelines/flux/modeling_flux.py | 3 --- src/diffusers/pipelines/ltx/modeling_latent_upsampler.py | 4 ---- src/diffusers/pipelines/ltx2/connectors.py | 7 ------- src/diffusers/pipelines/ltx2/latent_upsampler.py | 7 ------- .../pipelines/stable_audio/modeling_stable_audio.py | 5 ----- 7 files changed, 1 insertion(+), 34 deletions(-) diff --git a/src/diffusers/pipelines/ace_step/modeling_ace_step.py b/src/diffusers/pipelines/ace_step/modeling_ace_step.py index 26f6dcee3ffa..199e97fc9169 100644 --- a/src/diffusers/pipelines/ace_step/modeling_ace_step.py +++ b/src/diffusers/pipelines/ace_step/modeling_ace_step.py @@ -29,10 +29,7 @@ from ...models.condition_embedders.condition_embedder_ace_step import ( AceStepConditionEncoder as _AceStepConditionEncoder, ) -from ...models.condition_embedders.condition_embedder_ace_step import ( - AceStepEncoderLayer, - _pack_sequences, # noqa: F401 -) +from ...models.condition_embedders.condition_embedder_ace_step import AceStepEncoderLayer from ...models.condition_embedders.condition_embedder_ace_step import AceStepLyricEncoder as _AceStepLyricEncoder from ...models.condition_embedders.condition_embedder_ace_step import AceStepTimbreEncoder as _AceStepTimbreEncoder from ...models.modeling_utils import ModelMixin diff --git a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py index 9a30dc5be62c..6b918cadf800 100644 --- a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py @@ -30,10 +30,6 @@ from ...models.condition_embedders.condition_embedder_audioldm2 import ( AudioLDM2ProjectionModel as _AudioLDM2ProjectionModel, ) -from ...models.condition_embedders.condition_embedder_audioldm2 import ( - AudioLDM2ProjectionModelOutput, # noqa: F401 - add_special_tokens, # noqa: F401 -) from ...models.embeddings import TimestepEmbedding, Timesteps from ...models.modeling_utils import ModelMixin from ...models.resnet import Downsample2D, ResnetBlock2D, Upsample2D diff --git a/src/diffusers/pipelines/flux/modeling_flux.py b/src/diffusers/pipelines/flux/modeling_flux.py index 9e8647c3c0b3..2da7c2ce87b4 100644 --- a/src/diffusers/pipelines/flux/modeling_flux.py +++ b/src/diffusers/pipelines/flux/modeling_flux.py @@ -13,9 +13,6 @@ # limitations under the License. from ...models.condition_embedders.condition_embedder_flux_redux import ReduxImageEncoder as _ReduxImageEncoder -from ...models.condition_embedders.condition_embedder_flux_redux import ( - ReduxImageEncoderOutput, # noqa: F401 -) from ...utils import deprecate diff --git a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py index 67e0593278c0..58928a20cc19 100644 --- a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py +++ b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py @@ -13,10 +13,6 @@ # limitations under the License. from ...models.latent_upscaler.latent_upscaler_ltx import LTXLatentUpsamplerModel as _LTXLatentUpsamplerModel -from ...models.latent_upscaler.latent_upscaler_ltx import ( - PixelShuffleND, # noqa: F401 - ResBlock, # noqa: F401 -) from ...utils import deprecate diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py index 56034b2866da..5ac68d1ef846 100644 --- a/src/diffusers/pipelines/ltx2/connectors.py +++ b/src/diffusers/pipelines/ltx2/connectors.py @@ -12,13 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ...models.condition_embedders.condition_embedder_ltx2 import ( - LTX2ConnectorTransformer1d, # noqa: F401 - LTX2RotaryPosEmbed1d, # noqa: F401 - LTX2TransformerBlock1d, # noqa: F401 - per_layer_masked_mean_norm, # noqa: F401 - per_token_rms_norm, # noqa: F401 -) from ...models.condition_embedders.condition_embedder_ltx2 import LTX2TextConnectors as _LTX2TextConnectors from ...utils import deprecate diff --git a/src/diffusers/pipelines/ltx2/latent_upsampler.py b/src/diffusers/pipelines/ltx2/latent_upsampler.py index faf8583c04ad..b4a9972c2166 100644 --- a/src/diffusers/pipelines/ltx2/latent_upsampler.py +++ b/src/diffusers/pipelines/ltx2/latent_upsampler.py @@ -12,13 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ...models.latent_upscaler.latent_upscaler_ltx2 import ( - RATIONAL_RESAMPLER_SCALE_MAPPING, # noqa: F401 - BlurDownsample, # noqa: F401 - PixelShuffleND, # noqa: F401 - ResBlock, # noqa: F401 - SpatialRationalResampler, # noqa: F401 -) from ...models.latent_upscaler.latent_upscaler_ltx2 import LTX2LatentUpsamplerModel as _LTX2LatentUpsamplerModel from ...utils import deprecate diff --git a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py index 12fa39096912..17caa7e2a219 100644 --- a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py +++ b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py @@ -12,11 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ...models.condition_embedders.condition_embedder_stable_audio import ( - StableAudioNumberConditioner, # noqa: F401 - StableAudioPositionalEmbedding, # noqa: F401 - StableAudioProjectionModelOutput, # noqa: F401 -) from ...models.condition_embedders.condition_embedder_stable_audio import ( StableAudioProjectionModel as _StableAudioProjectionModel, ) From f52613e77074161701ae9bfb11081168d465229a Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 24 Oct 2026 04:06:19 +0000 Subject: [PATCH 5/5] helper --- src/diffusers/loaders/single_file.py | 9 +++-- .../pipelines/pipeline_loading_utils.py | 30 +++++------------ src/diffusers/utils/__init__.py | 7 +++- src/diffusers/utils/deprecation_utils.py | 33 +++++++++++++++++++ 4 files changed, 51 insertions(+), 28 deletions(-) diff --git a/src/diffusers/loaders/single_file.py b/src/diffusers/loaders/single_file.py index 91e5320e38e3..fe9565b0635c 100644 --- a/src/diffusers/loaders/single_file.py +++ b/src/diffusers/loaders/single_file.py @@ -21,7 +21,7 @@ from packaging import version from typing_extensions import Self -from ..utils import deprecate, is_transformers_available, logging +from ..utils import _maybe_remap_relocated_pipeline_class, deprecate, is_transformers_available, logging from .single_file_utils import ( SingleFileComponentError, _is_legacy_scheduler_kwargs, @@ -64,10 +64,9 @@ def load_single_file_sub_model( disable_mmap=False, **kwargs, ): - from ..pipelines.pipeline_loading_utils import _RELOCATED_PIPELINE_CLASSES - - if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: - library_name = "diffusers" + relocated_library_name = _maybe_remap_relocated_pipeline_class(library_name, class_name) + if relocated_library_name is not None: + library_name = relocated_library_name is_pipeline_module = False if is_pipeline_module: diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index c781d5df0100..2fe8051699fd 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -31,6 +31,7 @@ ONNX_WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, + _maybe_remap_relocated_pipeline_class, _maybe_remap_transformers_class, deprecate, get_class_from_dynamic_module, @@ -343,28 +344,13 @@ def _unwrap_model(model): return model -# Model classes that moved from a pipeline folder to `diffusers.models`. Existing `model_index.json` files still -# reference them as `[, ]`, so they are resolved to the top-level `diffusers` export -# instead of the deprecated alias left in the pipeline folder. -_RELOCATED_PIPELINE_CLASSES = { - ("ace_step", "AceStepConditionEncoder"), - ("audioldm2", "AudioLDM2ProjectionModel"), - ("flux", "ReduxImageEncoder"), - ("ltx", "LTXLatentUpsamplerModel"), - ("ltx2", "LTX2LatentUpsamplerModel"), - ("ltx2", "LTX2TextConnectors"), - ("stable_audio", "StableAudioProjectionModel"), - ("stable_audio_3", "StableAudio3DurationEmbedder"), - ("stable_diffusion", "CLIPImageProjection"), -} - - def maybe_raise_or_warn( library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module ): """Simple helper method to raise or warn in case incorrect module has been passed""" - if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: - library_name = "diffusers" + relocated_library_name = _maybe_remap_relocated_pipeline_class(library_name, class_name) + if relocated_library_name is not None: + library_name = relocated_library_name is_pipeline_module = False if not is_pipeline_module: @@ -401,8 +387,7 @@ def maybe_raise_or_warn( def simple_get_class_obj(library_name, class_name): from diffusers import pipelines - if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: - library_name = "diffusers" + library_name = _maybe_remap_relocated_pipeline_class(library_name, class_name) or library_name is_pipeline_module = hasattr(pipelines, library_name) @@ -437,8 +422,9 @@ def get_class_obj_and_candidates( if class_name.startswith("FlashPack"): class_name = class_name.removeprefix("FlashPack") - if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: - library_name = "diffusers" + relocated_library_name = _maybe_remap_relocated_pipeline_class(library_name, class_name) + if relocated_library_name is not None: + library_name = relocated_library_name is_pipeline_module = False if is_pipeline_module: diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index c25edb7f7c52..cf150eab2302 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -41,7 +41,12 @@ WEIGHTS_INDEX_NAME, WEIGHTS_NAME, ) -from .deprecation_utils import _maybe_remap_transformers_class, _resolve_dtype, deprecate +from .deprecation_utils import ( + _maybe_remap_relocated_pipeline_class, + _maybe_remap_transformers_class, + _resolve_dtype, + deprecate, +) from .doc_utils import replace_example_docstring from .dynamic_modules_utils import get_class_from_dynamic_module from .export_utils import encode_video, export_to_gif, export_to_obj, export_to_ply, export_to_video diff --git a/src/diffusers/utils/deprecation_utils.py b/src/diffusers/utils/deprecation_utils.py index 27cee39c73a4..48b427d9e1ff 100644 --- a/src/diffusers/utils/deprecation_utils.py +++ b/src/diffusers/utils/deprecation_utils.py @@ -68,6 +68,39 @@ def _maybe_remap_transformers_class(class_name: str) -> str | None: return None +# Model classes that moved from a pipeline folder to `diffusers.models`. Existing `model_index.json` files still +# reference them as `[, ]`, so they are resolved to the top-level `diffusers` export +# instead of the deprecated alias left in the pipeline folder. +_RELOCATED_PIPELINE_CLASSES = { + ("ace_step", "AceStepConditionEncoder"), + ("audioldm2", "AudioLDM2ProjectionModel"), + ("flux", "ReduxImageEncoder"), + ("ltx", "LTXLatentUpsamplerModel"), + ("ltx2", "LTX2LatentUpsamplerModel"), + ("ltx2", "LTX2TextConnectors"), + ("stable_audio", "StableAudioProjectionModel"), + ("stable_audio_3", "StableAudio3DurationEmbedder"), + ("stable_diffusion", "CLIPImageProjection"), +} + + +def _maybe_remap_relocated_pipeline_class(library_name: str, class_name: str) -> str | None: + """ + Check if a class listed under a pipeline folder has moved to `diffusers.models`. + + Args: + library_name: The library or pipeline folder the class is listed under + class_name: The name of the class to check + + Returns: + The library to load the class from if it has moved, None otherwise + """ + if (library_name, class_name) in _RELOCATED_PIPELINE_CLASSES: + return "diffusers" + + return None + + def deprecate(*args, take_from: dict | Any | None = None, standard_warn=True, stacklevel=2): from .. import __version__