import torch
from diffusers import StableAudioPipeline, StableAudioProjectionModel, StableAudioDiTModel, EDMDPMSolverMultistepScheduler, AutoencoderOobleck
from transformers import T5EncoderModel, T5Tokenizer

def get_dummy_components():
    torch.manual_seed(0)
    transformer = StableAudioDiTModel(
        sample_size=32,
        in_channels=6,
        num_layers=2,
        attention_head_dim=4,
        num_key_value_attention_heads=2,
        out_channels=6,
        cross_attention_dim=4,
        time_proj_dim=8,
        global_states_input_dim=48,
        cross_attention_input_dim=24
    )
    scheduler = EDMDPMSolverMultistepScheduler(
        solver_order=2,
        prediction_type="v_prediction",
        noise_preconditioning_strategy="atan",
        sigma_data=1.0,
        algorithm_type="sde-dpmsolver++",
        sigma_schedule="exponential",
        noise_sampling_strategy="brownian_tree",
    )
    torch.manual_seed(0)
    vae = AutoencoderOobleck(
        encoder_hidden_size=12,
        downsampling_ratios=[1, 2],
        decoder_channels=12,
        decoder_input_channels=6,
        audio_channels=2,
        channel_multiples=[2, 4],
        sampling_rate=32,
    )
    torch.manual_seed(0)
    t5_repo_id = "hf-internal-testing/tiny-random-T5ForConditionalGeneration"
    text_encoder = T5EncoderModel.from_pretrained(t5_repo_id)
    tokenizer = T5Tokenizer.from_pretrained(t5_repo_id, truncation=True, model_max_length=25)

    torch.manual_seed(0)
    projection_model = StableAudioProjectionModel(
        text_encoder_dim=text_encoder.config.d_model,
        conditioning_dim=24,
        min_value=0,
        max_value=256,
    )
    components = {
        "transformer": transformer,
        "scheduler": scheduler,
        "vae": vae,
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "projection_model": projection_model,
    }
    return components


def get_dummy_inputs():
    generator = torch.manual_seed(0)
    inputs = {
        "prompt": "A hammer hitting a wooden surface",
        "generator": generator,
        "num_inference_steps": 2,
        "guidance_scale": 6.0,
    }
    return inputs

components = get_dummy_components()
pipeline = StableAudioPipeline(**components)
pipeline.enable_sequential_cpu_offload(device="cuda")

inputs = get_dummy_inputs()
pipeline(**inputs)