import torch
from diffusers import AudioLDM2ProjectionModel

model = AudioLDM2ProjectionModel(text_encoder_dim=3, text_encoder_1_dim=4, langauge_model_dim=5)
h0 = torch.randn(2, 1, 3)
h1 = torch.randn(2, 7, 4)
mask1 = torch.ones(2, 7, dtype=torch.long)

model(hidden_states=h0, hidden_states_1=h1, attention_mask=None, attention_mask_1=mask1)
# TypeError: new_ones(): argument 'size' must be tuple of ints

# before concatenating hidden_states and hidden_states_1
if attention_mask is None and attention_mask_1 is not None:
    attention_mask = attention_mask_1.new_ones(hidden_states.shape[:2])
elif attention_mask is not None and attention_mask_1 is None:
    attention_mask_1 = attention_mask.new_ones(hidden_states_1.shape[:2])

hidden_states = torch.cat([hidden_states, hidden_states_1], dim=1)

import torch
from diffusers import AudioLDM2UNet2DConditionModel

torch.manual_seed(0)
model = AudioLDM2UNet2DConditionModel(
    sample_size=8, in_channels=4, out_channels=4, block_out_channels=(8,),
    layers_per_block=1, norm_num_groups=1,
    down_block_types=("CrossAttnDownBlock2D",),
    up_block_types=("CrossAttnUpBlock2D",),
    cross_attention_dim=((None, 8, 8),),
    attention_head_dim=1,
).eval()

sample = torch.randn(1, 4, 8, 8)
encoder = torch.randn(1, 5, 8)
mask = torch.tensor([[1, 1, 1, 0, 0]])

with torch.no_grad():
    implicit = model(sample, 1, encoder_hidden_states=encoder, encoder_attention_mask=mask).sample
    explicit = model(
        sample, 1, encoder_hidden_states=encoder, encoder_attention_mask=mask,
        encoder_hidden_states_1=encoder, encoder_attention_mask_1=mask,
    ).sample

print((implicit - explicit).abs().max())
# tensor(0.1561...)

if encoder_hidden_states_1 is None:
    encoder_hidden_states_1 = encoder_hidden_states
    if encoder_attention_mask_1 is None:
        encoder_attention_mask_1 = encoder_attention_mask

import torch
from diffusers import AudioLDM2UNet2DConditionModel
from diffusers.models.attention_processor import AttnProcessor

class MarkerProcessor(AttnProcessor):
    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None, marker=None, **kwargs):
        if marker:
            raise RuntimeError("cross_attention_kwargs propagated")
        return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask, temb, **kwargs)

model = AudioLDM2UNet2DConditionModel(
    sample_size=8, in_channels=4, out_channels=4, block_out_channels=(8,),
    layers_per_block=1, norm_num_groups=1,
    down_block_types=("CrossAttnDownBlock2D",),
    up_block_types=("CrossAttnUpBlock2D",),
    cross_attention_dim=(8,),
    attention_head_dim=1,
).eval()
model.set_attn_processor(MarkerProcessor())

with torch.no_grad():
    model(torch.randn(1, 4, 8, 8), 1, encoder_hidden_states=torch.randn(1, 5, 8), cross_attention_kwargs={"marker": True})

print("marker was ignored")

# pipeline denoise call
noise_pred = self.unet(
    latent_model_input,
    t,
    encoder_hidden_states=generated_prompt_embeds,
    encoder_hidden_states_1=prompt_embeds,
    encoder_attention_mask_1=attention_mask,
    cross_attention_kwargs=cross_attention_kwargs,
    return_dict=False,
)[0]

# all AudioLDM2 cross-attn block normal paths
hidden_states = self.attentions[i * num_attention_per_layer + idx](
    hidden_states,
    encoder_hidden_states=forward_encoder_hidden_states,
    cross_attention_kwargs=cross_attention_kwargs,
    attention_mask=attention_mask,
    encoder_attention_mask=forward_encoder_attention_mask,
    return_dict=False,
)[0]

import torch
from diffusers import AudioLDM2UNet2DConditionModel

model = AudioLDM2UNet2DConditionModel(
    sample_size=8, in_channels=4, out_channels=4, block_out_channels=(8,),
    layers_per_block=1, norm_num_groups=1,
    down_block_types=("CrossAttnDownBlock2D",),
    up_block_types=("CrossAttnUpBlock2D",),
    cross_attention_dim=(8,),
    attention_head_dim=1,
)
model.enable_gradient_checkpointing()

sample = torch.randn(1, 4, 8, 8, requires_grad=True)
encoder = torch.randn(1, 5, 8, requires_grad=True)
mask = torch.tensor([[1, 1, 1, 0, 0]])

model(sample, 1, encoder_hidden_states=encoder, encoder_attention_mask=mask)
# RuntimeError: The size of tensor a (64) must match the size of tensor b (69)...

hidden_states = self._gradient_checkpointing_func(self.resnets[i], hidden_states, temb)
hidden_states = self.attentions[i * num_attention_per_layer + idx](
    hidden_states,
    encoder_hidden_states=forward_encoder_hidden_states,
    cross_attention_kwargs=cross_attention_kwargs,
    attention_mask=attention_mask,
    encoder_attention_mask=forward_encoder_attention_mask,
    return_dict=False,
)[0]

import torch
from types import SimpleNamespace
from diffusers import AudioLDM2Pipeline

class Inputs(dict):
    def to(self, device): return self

class Tokenizer:
    def __call__(self, text, return_tensors=None, padding=None):
        return Inputs(input_ids=torch.ones(len(text), 2, dtype=torch.long))

class FeatureExtractor:
    sampling_rate = 16000
    def __call__(self, audio, return_tensors=None, sampling_rate=None):
        return SimpleNamespace(input_features=torch.zeros(len(audio), 1, 4))

class TextEncoder:
    def __call__(self, **inputs):
        return SimpleNamespace(logits_per_text=torch.tensor([[0.1, 0.2, 0.9, 0.8], [0.7, 0.6, 0.5, 0.4]]))

pipe = AudioLDM2Pipeline.__new__(AudioLDM2Pipeline)
pipe.tokenizer = Tokenizer()
pipe.feature_extractor = FeatureExtractor()
pipe.text_encoder = TextEncoder()
pipe.vocoder = SimpleNamespace(config=SimpleNamespace(sampling_rate=16000))

audio = torch.arange(16, dtype=torch.float32).view(4, 4)
ranked = pipe.score_waveforms(["prompt a", "prompt b"], audio, num_waveforms_per_prompt=2, device="cpu", dtype=torch.float32)
print(ranked[:, 0].tolist())
# [8.0, 12.0, 0.0, 4.0] selects rows 2/3 for prompt 0 and rows 0/1 for prompt 1

batch_size = len(text) if isinstance(text, list) else 1
selected = []
for i in range(batch_size):
    start = i * num_waveforms_per_prompt
    end = start + num_waveforms_per_prompt
    local_scores = logits_per_text[i, start:end]
    local_indices = torch.argsort(local_scores, descending=True) + start
    selected.append(local_indices)

indices = torch.cat(selected)
audio = torch.index_select(audio, 0, indices.cpu())

from diffusers import AudioLDM2UNet2DConditionModel

AudioLDM2UNet2DConditionModel(
    block_out_channels=(8, 16),
    layers_per_block=1,
    norm_num_groups=1,
    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
    cross_attention_dim=(8,),
)
# IndexError: tuple index out of range

if not isinstance(cross_attention_dim, int) and len(cross_attention_dim) != len(down_block_types):
    raise ValueError(
        f"Must provide the same number of `cross_attention_dim` as `down_block_types`. "
        f"`cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
    )

if not isinstance(transformer_layers_per_block, int) and len(transformer_layers_per_block) != len(down_block_types):
    raise ValueError(
        f"Must provide the same number of `transformer_layers_per_block` as `down_block_types`. "
        f"`transformer_layers_per_block`: {transformer_layers_per_block}. `down_block_types`: {down_block_types}."
    )