upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

Author: RuixiangMa

docs/source/en/api/pipelines/longcat_audio_dit.md

@@ -26,34 +26,32 @@ This pipeline was adapted from the LongCat-AudioDiT reference implementation: ht
 ## Usage
 
 ```py
+import soundfile as sf
 import torch
 from diffusers import LongCatAudioDiTPipeline
 
-repo_id = "<longcat-audio-dit-repo-id>"
-tokenizer_path = os.environ["LONGCAT_AUDIO_DIT_TOKENIZER_PATH"]
-
-pipe = LongCatAudioDiTPipeline.from_pretrained(
-    repo_id,
-    tokenizer=tokenizer_path,
+pipeline = LongCatAudioDiTPipeline.from_pretrained(
+    "meituan-longcat/LongCat-AudioDiT-1B",
     torch_dtype=torch.float16,
-    local_files_only=True,
 )
-pipe = pipe.to("cuda")
+pipeline = pipeline.to("cuda")
 
-audio = pipe(
+audio = pipeline(
     prompt="A calm ocean wave ambience with soft wind in the background.",
-    audio_end_in_s=2.0,
+    audio_end_in_s=5.0,
     num_inference_steps=16,
     guidance_scale=4.0,
     output_type="pt",
 ).audios
+
+output = audio[0, 0].float().cpu().numpy()
+sf.write("longcat.wav", output, pipeline.sample_rate)
 ```
 
 ## Tips
 
 - `audio_end_in_s` is the most direct way to control output duration.
 - `output_type="pt"` returns a PyTorch tensor shaped `(batch, channels, samples)`.
-- If your tokenizer path is local-only, pass it explicitly to `from_pretrained(...)`.
 
 ## LongCatAudioDiTPipeline
 

src/diffusers/pipelines/longcat_audio_dit/pipeline_longcat_audio_dit.py

@@ -26,7 +26,7 @@
 from huggingface_hub.utils import validate_hf_hub_args
 from safetensors.torch import load_file
 from torch.nn.utils.rnn import pad_sequence
-from transformers import PreTrainedTokenizerBase, T5Tokenizer, UMT5Config, UMT5EncoderModel
+from transformers import AutoTokenizer, PreTrainedTokenizerBase, UMT5Config, UMT5EncoderModel
 
 from ...models import LongCatAudioDiTTransformer, LongCatAudioDiTVae
 from ...utils import HUGGINGFACE_CO_RESOLVE_ENDPOINT, logging
@@ -105,7 +105,7 @@ def _load_longcat_tokenizer(
     tokenizer_kwargs = {"local_files_only": local_files_only}
     if not isinstance(tokenizer_source, Path) and tokenizer_source == pretrained_model_name_or_path and subfolder:
         tokenizer_kwargs["subfolder"] = subfolder
-    return T5Tokenizer.from_pretrained(tokenizer_source, **tokenizer_kwargs)
+    return AutoTokenizer.from_pretrained(tokenizer_source, **tokenizer_kwargs)
 
 
 def _resolve_longcat_file(
@@ -278,6 +278,10 @@ def from_pretrained(
             transformer = transformer.to(dtype=torch_dtype)
             vae = vae.to(dtype=torch_dtype)
 
+        text_encoder.eval()
+        transformer.eval()
+        vae.eval()
+
         pipe = cls(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=transformer)
         pipe.sample_rate = config.get("sampling_rate", pipe.sample_rate)
         pipe.latent_hop = config.get("latent_hop", pipe.latent_hop)
@@ -322,15 +326,24 @@ def prepare_latents(
         dtype: torch.dtype,
         generator: torch.Generator | list[torch.Generator] | None = None,
     ) -> torch.Tensor:
+        if isinstance(generator, list):
+            if len(generator) != batch_size:
+                raise ValueError(
+                    f"Expected {batch_size} generators for batch size {batch_size}, but got {len(generator)}."
+                )
+            generators = generator
+        else:
+            generators = [generator] * batch_size
+
         latents = [
             torch.randn(
                 duration,
                 self.latent_dim,
                 device=device,
                 dtype=dtype,
-                generator=generator if isinstance(generator, torch.Generator) else None,
+                generator=generators[idx],
             )
-            for _ in range(batch_size)
+            for idx in range(batch_size)
         ]
         return pad_sequence(latents, padding_value=0.0, batch_first=True)
 
@@ -381,6 +394,12 @@ def __call__(
         else:
             if isinstance(negative_prompt, str):
                 negative_prompt = [negative_prompt] * batch_size
+            else:
+                negative_prompt = list(negative_prompt)
+                if len(negative_prompt) != batch_size:
+                    raise ValueError(
+                        f"`negative_prompt` must have batch size {batch_size}, but got {len(negative_prompt)} prompts."
+                    )
             neg_text, neg_text_len = self.encode_prompt(negative_prompt, device)
             neg_text_mask = _lens_to_mask(neg_text_len, length=neg_text.shape[1])
 
@@ -399,7 +418,7 @@ def model_step(curr_t: torch.Tensor, current_sample: torch.Tensor) -> torch.Tens
                 attention_mask=mask,
                 latent_cond=latent_cond,
             ).sample
-            if guidance_scale < 1e-5:
+            if guidance_scale <= 1.0:
                 return pred
             null_pred = self.transformer(
                 hidden_states=current_sample,
@@ -409,7 +428,7 @@ def model_step(curr_t: torch.Tensor, current_sample: torch.Tensor) -> torch.Tens
                 attention_mask=mask,
                 latent_cond=latent_cond,
             ).sample
-            return pred + (pred - null_pred) * guidance_scale
+            return null_pred + (pred - null_pred) * guidance_scale
 
         for idx in range(len(timesteps) - 1):
             curr_t = timesteps[idx]