import torch
from types import SimpleNamespace
from diffusers import LongCatAudioDiTPipeline, LongCatAudioDiTTransformer, LongCatAudioDiTVae

class Tok:
    model_max_length = 4
    def __call__(self, prompt, **kwargs):
        b = len(prompt)
        return SimpleNamespace(input_ids=torch.ones(b, 4, dtype=torch.long), attention_mask=torch.ones(b, 4))

class Text(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = torch.nn.Embedding(8, 32)
    def forward(self, input_ids, attention_mask, output_hidden_states=False):
        h = self.emb(input_ids)  # float32
        return SimpleNamespace(last_hidden_state=h, hidden_states=(h,))

pipe = LongCatAudioDiTPipeline(
    vae=LongCatAudioDiTVae(in_channels=1, channels=4, c_mults=[1], strides=[2], latent_dim=8, encoder_latent_dim=16),
    text_encoder=Text(),
    tokenizer=Tok(),
    transformer=LongCatAudioDiTTransformer(dit_dim=64, dit_depth=1, dit_heads=4, dit_text_dim=32, latent_dim=8, text_conv=False).to(torch.bfloat16),
)
pipe.set_progress_bar_config(disable=True)
pipe("x", audio_duration_s=0.1, num_inference_steps=1, guidance_scale=1.0, output_type="latent")

latent_dtype = self.transformer.dtype

prompt_embeds = prompt_embeds.to(dtype=latent_dtype)
negative_prompt_embeds = negative_prompt_embeds.to(dtype=latent_dtype)

latent_cond = torch.zeros(batch_size, duration, self.latent_dim, device=device, dtype=latent_dtype)
latents = self.prepare_latents(batch_size, duration, device, latent_dtype, generator=generator, latents=latents)

curr_t = (t / self.scheduler.config.num_train_timesteps).expand(batch_size).to(dtype=latent_dtype)

import torch
from types import SimpleNamespace
from diffusers import LongCatAudioDiTPipeline, LongCatAudioDiTTransformer, LongCatAudioDiTVae

class Tok:
    model_max_length = 4
    def __call__(self, prompt, **kwargs):
        return SimpleNamespace(input_ids=torch.ones(1, 4, dtype=torch.long), attention_mask=torch.ones(1, 4))

class Text(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = torch.nn.Embedding(8, 32)
    def forward(self, input_ids, attention_mask, output_hidden_states=False):
        h = self.emb(input_ids)
        return SimpleNamespace(last_hidden_state=h, hidden_states=(h,))

pipe = LongCatAudioDiTPipeline(
    vae=LongCatAudioDiTVae(in_channels=1, channels=4, c_mults=[1], strides=[2], latent_dim=8, encoder_latent_dim=16),
    text_encoder=Text(),
    tokenizer=Tok(),
    transformer=LongCatAudioDiTTransformer(dit_dim=64, dit_depth=1, dit_heads=4, dit_text_dim=32, latent_dim=8, text_conv=False),
)
with torch.enable_grad():
    embeds, _ = pipe.encode_prompt("x", torch.device("cpu"))
print(embeds.requires_grad, embeds.grad_fn)  # False, None

# Remove the inner no_grad block.
output = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)

import torch
from diffusers import LongCatAudioDiTVae

vae = LongCatAudioDiTVae(
    in_channels=1, channels=16, c_mults=[1, 2], strides=[2],
    latent_dim=8, encoder_latent_dim=16, downsampling_ratio=2,
)
decoded = vae.decode(torch.zeros(1, 8, 10)).sample
print("configured ratio:", vae.config.downsampling_ratio)
print("actual ratio:", decoded.shape[-1] // 10)  # 4

actual_downsampling_ratio = math.prod(strides)
if downsampling_ratio != actual_downsampling_ratio:
    raise ValueError(
        f"`downsampling_ratio` ({downsampling_ratio}) must match product(strides) ({actual_downsampling_ratio})."
    )

from diffusers import LongCatAudioDiTTransformer

model = LongCatAudioDiTTransformer(dit_dim=64, dit_depth=1, dit_heads=4, dit_text_dim=32, latent_dim=8, text_conv=False)
print(hasattr(model, "attn_processors"))  # False

try:
    model.blocks[0].self_attn.fuse_projections()
except Exception as e:
    print(type(e).__name__, e)  # AttributeError: no attribute 'use_bias'

from ..attention import AttentionMixin, AttentionModuleMixin

class AudioDiTAttention(nn.Module, AttentionModuleMixin):
    _default_processor_cls = AudioDiTSelfAttnProcessor
    _available_processors = [AudioDiTSelfAttnProcessor, AudioDiTCrossAttnProcessor]
    _supports_qkv_fusion = False

    def __init__(..., bias: bool = True, ...):
        ...
        self.use_bias = bias
        self.set_processor(processor or self._default_processor_cls())

class LongCatAudioDiTTransformer(ModelMixin, AttentionMixin, ConfigMixin):
    ...

import tempfile
from pathlib import Path
from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel
from diffusers import LongCatAudioDiTPipeline, LongCatAudioDiTTransformer, LongCatAudioDiTVae

tok = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
pipe = LongCatAudioDiTPipeline(
    vae=LongCatAudioDiTVae(in_channels=1, channels=4, c_mults=[1], strides=[2], latent_dim=8, encoder_latent_dim=16),
    text_encoder=UMT5EncoderModel(UMT5Config(d_model=32, num_layers=1, num_heads=4, d_ff=64, vocab_size=tok.vocab_size)),
    tokenizer=tok,
    transformer=LongCatAudioDiTTransformer(dit_dim=64, dit_depth=1, dit_heads=4, dit_text_dim=32, latent_dim=8, text_conv=False),
)
with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as d:
    pipe.save_pretrained(d)
    LongCatAudioDiTPipeline.from_pretrained(d, tokenizer=Path(d) / "tokenizer", local_files_only=True)

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
pipe = LongCatAudioDiTPipeline.from_pretrained(
    model_path,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    local_files_only=True,
)