Address ACE-Step audio token and APG review

Author: ChuxiJ

scripts/convert_ace_step_to_diffusers.py

@@ -131,7 +131,9 @@ def convert_ace_step_weights(checkpoint_dir, dit_config, output_dir, dtype_str="
     # =========================================================================
     transformer_sd = {}
     condition_encoder_sd = {}
-    other_sd = {}  # tokenizer, detokenizer (audio quantization — not used by the text2music pipeline)
+    audio_tokenizer_sd = {}
+    audio_token_detokenizer_sd = {}
+    other_sd = {}
 
     # Rename original ACE-Step attention keys to the diffusers `Attention` +
     # `AttnProcessor` convention (`to_q`/`to_k`/`to_v`/`to_out.0`/`norm_q`/`norm_k`).
@@ -174,11 +176,21 @@ def _rename_attn_keys(key: str) -> str:
             # Keep it co-located with the condition encoder since that is where the
             # pipeline pulls unconditional sequences from.
             condition_encoder_sd["null_condition_emb"] = value.to(target_dtype)
+        elif key.startswith("tokenizer."):
+            new_key = key[len("tokenizer.") :]
+            new_key = _rename_attn_keys(new_key)
+            audio_tokenizer_sd[new_key] = value.to(target_dtype)
+        elif key.startswith("detokenizer."):
+            new_key = key[len("detokenizer.") :]
+            new_key = _rename_attn_keys(new_key)
+            audio_token_detokenizer_sd[new_key] = value.to(target_dtype)
         else:
             other_sd[key] = value.to(target_dtype)
 
     print(f"  Transformer keys: {len(transformer_sd)}")
     print(f"  Condition encoder keys: {len(condition_encoder_sd)}")
+    print(f"  Audio tokenizer keys: {len(audio_tokenizer_sd)}")
+    print(f"  Audio token detokenizer keys: {len(audio_token_detokenizer_sd)}")
     print(f"  Other keys: {len(other_sd)} ({list(other_sd.keys())[:5]}...)")
 
     # =========================================================================
@@ -248,6 +260,47 @@ def _rename_attn_keys(key: str) -> str:
         "sliding_window": original_config["sliding_window"],
     }
 
+    audio_tokenizer_config = {
+        "_class_name": "AceStepAudioTokenizer",
+        "_diffusers_version": "0.33.0.dev0",
+        "hidden_size": encoder_hidden_size,
+        "intermediate_size": encoder_intermediate_size,
+        "audio_acoustic_hidden_dim": original_config["audio_acoustic_hidden_dim"],
+        "pool_window_size": original_config.get("pool_window_size", 5),
+        "fsq_dim": original_config.get("fsq_dim", encoder_hidden_size),
+        "fsq_input_levels": original_config.get("fsq_input_levels", [8, 8, 8, 5, 5, 5]),
+        "fsq_input_num_quantizers": original_config.get("fsq_input_num_quantizers", 1),
+        "num_attention_pooler_hidden_layers": original_config.get("num_attention_pooler_hidden_layers", 2),
+        "num_attention_heads": encoder_num_attention_heads,
+        "num_key_value_heads": encoder_num_key_value_heads,
+        "head_dim": original_config["head_dim"],
+        "rope_theta": original_config["rope_theta"],
+        "attention_bias": original_config["attention_bias"],
+        "attention_dropout": original_config["attention_dropout"],
+        "rms_norm_eps": original_config["rms_norm_eps"],
+        "sliding_window": original_config["sliding_window"],
+        "layer_types": original_config["layer_types"][: original_config.get("num_attention_pooler_hidden_layers", 2)],
+    }
+
+    audio_token_detokenizer_config = {
+        "_class_name": "AceStepAudioTokenDetokenizer",
+        "_diffusers_version": "0.33.0.dev0",
+        "hidden_size": encoder_hidden_size,
+        "intermediate_size": encoder_intermediate_size,
+        "audio_acoustic_hidden_dim": original_config["audio_acoustic_hidden_dim"],
+        "pool_window_size": original_config.get("pool_window_size", 5),
+        "num_attention_pooler_hidden_layers": original_config.get("num_attention_pooler_hidden_layers", 2),
+        "num_attention_heads": encoder_num_attention_heads,
+        "num_key_value_heads": encoder_num_key_value_heads,
+        "head_dim": original_config["head_dim"],
+        "rope_theta": original_config["rope_theta"],
+        "attention_bias": original_config["attention_bias"],
+        "attention_dropout": original_config["attention_dropout"],
+        "rms_norm_eps": original_config["rms_norm_eps"],
+        "sliding_window": original_config["sliding_window"],
+        "layer_types": original_config["layer_types"][: original_config.get("num_attention_pooler_hidden_layers", 2)],
+    }
+
     # =========================================================================
     # 3. Bake silence_latent into the condition_encoder state dict.
     #
@@ -282,11 +335,19 @@ def _rename_attn_keys(key: str) -> str:
         AutoencoderOobleck,
         FlowMatchEulerDiscreteScheduler,
     )
-    from diffusers.pipelines.ace_step import AceStepConditionEncoder
+    from diffusers.pipelines.ace_step import (
+        AceStepAudioTokenDetokenizer,
+        AceStepAudioTokenizer,
+        AceStepConditionEncoder,
+    )
 
     # Drop metadata keys — they're re-populated by `save_pretrained` at save time.
     transformer_init_kwargs = {k: v for k, v in transformer_config.items() if not k.startswith("_")}
     condition_encoder_init_kwargs = {k: v for k, v in condition_encoder_config.items() if not k.startswith("_")}
+    audio_tokenizer_init_kwargs = {k: v for k, v in audio_tokenizer_config.items() if not k.startswith("_")}
+    audio_token_detokenizer_init_kwargs = {
+        k: v for k, v in audio_token_detokenizer_config.items() if not k.startswith("_")
+    }
 
     print("\nConstructing transformer ...")
     transformer = AceStepTransformer1DModel(**transformer_init_kwargs).to(target_dtype)
@@ -296,6 +357,14 @@ def _rename_attn_keys(key: str) -> str:
     condition_encoder = AceStepConditionEncoder(**condition_encoder_init_kwargs).to(target_dtype)
     condition_encoder.load_state_dict(condition_encoder_sd, strict=True)
 
+    print("Constructing audio_tokenizer ...")
+    audio_tokenizer = AceStepAudioTokenizer(**audio_tokenizer_init_kwargs).to(target_dtype)
+    audio_tokenizer.load_state_dict(audio_tokenizer_sd, strict=True)
+
+    print("Constructing audio_token_detokenizer ...")
+    audio_token_detokenizer = AceStepAudioTokenDetokenizer(**audio_token_detokenizer_init_kwargs).to(target_dtype)
+    audio_token_detokenizer.load_state_dict(audio_token_detokenizer_sd, strict=True)
+
     print("Loading VAE ...")
     vae = AutoencoderOobleck.from_pretrained(vae_dir).to(target_dtype)
 
@@ -319,6 +388,8 @@ def _rename_attn_keys(key: str) -> str:
         transformer=transformer,
         condition_encoder=condition_encoder,
         scheduler=scheduler,
+        audio_tokenizer=audio_tokenizer,
+        audio_token_detokenizer=audio_token_detokenizer,
     )
 
     print(f"\nSaving pipeline -> {output_dir}")
@@ -331,18 +402,13 @@ def _rename_attn_keys(key: str) -> str:
         shutil.copy2(silence_latent_src, os.path.join(output_dir, "silence_latent.pt"))
         print(f"  kept raw silence_latent copy at {output_dir}/silence_latent.pt")
 
-    # Report other keys that were not saved to transformer or condition_encoder
+    # Report any keys that were not saved to registered pipeline modules.
     if other_sd:
-        print(f"\nNote: {len(other_sd)} keys were dropped (tokenizer / detokenizer weights):")
+        print(f"\nNote: {len(other_sd)} keys were dropped:")
         for key in sorted(other_sd.keys())[:10]:
             print(f"  {key}")
         if len(other_sd) > 10:
             print(f"  ... ({len(other_sd) - 10} more)")
-        print(
-            "These belong to the audio tokenizer / detokenizer used by the 5Hz LM path "
-            "(cover / audio-code tasks). The Diffusers text2music pipeline does not "
-            "currently expose them."
-        )
 
     print(f"\nConversion complete! Output saved to: {output_dir}")
     print("\nTo load the pipeline:")

src/diffusers/__init__.py

@@ -487,6 +487,8 @@
     )
     _import_structure["pipelines"].extend(
         [
+            "AceStepAudioTokenDetokenizer",
+            "AceStepAudioTokenizer",
             "AceStepConditionEncoder",
             "AceStepPipeline",
             "AllegroPipeline",
@@ -1277,6 +1279,8 @@
             ZImageModularPipeline,
         )
         from .pipelines import (
+            AceStepAudioTokenDetokenizer,
+            AceStepAudioTokenizer,
             AceStepConditionEncoder,
             AceStepPipeline,
             AllegroPipeline,

src/diffusers/guiders/adaptive_projected_guidance.py

@@ -40,6 +40,9 @@ class AdaptiveProjectedGuidance(BaseGuidance):
             The momentum parameter for the adaptive projected guidance. Disabled if set to `None`.
         adaptive_projected_guidance_rescale (`float`, defaults to `15.0`):
             The rescale factor applied to the noise predictions. This is used to improve image quality and fix
+        adaptive_projected_guidance_norm_dim (`int` or `tuple[int]`, *optional*):
+            Dimension(s) over which to compute the APG norm and projection. If omitted, all non-batch dimensions are
+            used, preserving the original behavior.
         guidance_rescale (`float`, defaults to `0.0`):
             The rescale factor applied to the noise predictions. This is used to improve image quality and fix
             overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
@@ -62,6 +65,7 @@ def __init__(
         guidance_scale: float = 7.5,
         adaptive_projected_guidance_momentum: float | None = None,
         adaptive_projected_guidance_rescale: float = 15.0,
+        adaptive_projected_guidance_norm_dim: int | tuple[int, ...] | None = None,
         eta: float = 1.0,
         guidance_rescale: float = 0.0,
         use_original_formulation: bool = False,
@@ -74,6 +78,7 @@ def __init__(
         self.guidance_scale = guidance_scale
         self.adaptive_projected_guidance_momentum = adaptive_projected_guidance_momentum
         self.adaptive_projected_guidance_rescale = adaptive_projected_guidance_rescale
+        self.adaptive_projected_guidance_norm_dim = adaptive_projected_guidance_norm_dim
         self.eta = eta
         self.guidance_rescale = guidance_rescale
         self.use_original_formulation = use_original_formulation
@@ -117,6 +122,7 @@ def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = No
                 self.eta,
                 self.adaptive_projected_guidance_rescale,
                 self.use_original_formulation,
+                self.adaptive_projected_guidance_norm_dim,
             )
 
         if self.guidance_rescale > 0.0:
@@ -210,9 +216,15 @@ def normalized_guidance(
     eta: float = 1.0,
     norm_threshold: float = 0.0,
     use_original_formulation: bool = False,
+    norm_dim: int | tuple[int, ...] | None = None,
 ):
     diff = pred_cond - pred_uncond
-    dim = [-i for i in range(1, len(diff.shape))]
+    if norm_dim is None:
+        dim = [-i for i in range(1, len(diff.shape))]
+    elif isinstance(norm_dim, int):
+        dim = [norm_dim]
+    else:
+        dim = list(norm_dim)
 
     if momentum_buffer is not None:
         momentum_buffer.update(diff)
@@ -224,11 +236,15 @@ def normalized_guidance(
         scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
         diff = diff * scale_factor
 
-    v0, v1 = diff.double(), pred_cond.double()
+    if diff.device.type in {"mps", "npu"}:
+        v0, v1 = diff.cpu().double(), pred_cond.cpu().double()
+    else:
+        v0, v1 = diff.double(), pred_cond.double()
     v1 = torch.nn.functional.normalize(v1, dim=dim)
     v0_parallel = (v0 * v1).sum(dim=dim, keepdim=True) * v1
     v0_orthogonal = v0 - v0_parallel
-    diff_parallel, diff_orthogonal = v0_parallel.type_as(diff), v0_orthogonal.type_as(diff)
+    diff_parallel = v0_parallel.to(device=diff.device, dtype=diff.dtype)
+    diff_orthogonal = v0_orthogonal.to(device=diff.device, dtype=diff.dtype)
     normalized_update = diff_orthogonal + eta * diff_parallel
 
     pred = pred_cond if use_original_formulation else pred_uncond

src/diffusers/pipelines/__init__.py

@@ -150,6 +150,8 @@
         ]
     )
     _import_structure["ace_step"] = [
+        "AceStepAudioTokenDetokenizer",
+        "AceStepAudioTokenizer",
         "AceStepConditionEncoder",
         "AceStepPipeline",
     ]
@@ -578,7 +580,12 @@
     except OptionalDependencyNotAvailable:
         from ..utils.dummy_torch_and_transformers_objects import *
     else:
-        from .ace_step import AceStepConditionEncoder, AceStepPipeline
+        from .ace_step import (
+            AceStepAudioTokenDetokenizer,
+            AceStepAudioTokenizer,
+            AceStepConditionEncoder,
+            AceStepPipeline,
+        )
         from .allegro import AllegroPipeline
         from .animatediff import (
             AnimateDiffControlNetPipeline,

src/diffusers/pipelines/ace_step/__init__.py

@@ -22,7 +22,11 @@
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
-    _import_structure["modeling_ace_step"] = ["AceStepConditionEncoder"]
+    _import_structure["modeling_ace_step"] = [
+        "AceStepAudioTokenDetokenizer",
+        "AceStepAudioTokenizer",
+        "AceStepConditionEncoder",
+    ]
     _import_structure["pipeline_ace_step"] = ["AceStepPipeline"]
 
 
@@ -34,7 +38,7 @@
         from ...utils.dummy_torch_and_transformers_objects import *
 
     else:
-        from .modeling_ace_step import AceStepConditionEncoder
+        from .modeling_ace_step import AceStepAudioTokenDetokenizer, AceStepAudioTokenizer, AceStepConditionEncoder
         from .pipeline_ace_step import AceStepPipeline
 
 else: