add files used by space

MeanAudio.py

@@ -0,0 +1,147 @@
+import warnings
+warnings.filterwarnings("ignore", category=FutureWarning)
+import logging
+from pathlib import Path
+import torch
+import torchaudio
+from meanaudio.eval_utils import (ModelConfig, all_model_cfg, generate_mf, generate_fm, setup_eval_logging)
+from meanaudio.model.flow_matching import FlowMatching
+from meanaudio.model.mean_flow import MeanFlow
+from meanaudio.model.networks import MeanAudio, get_mean_audio
+from meanaudio.model.utils.features_utils import FeaturesUtils
+from huggingface_hub import snapshot_download
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+log = logging.getLogger()
+
+@torch.inference_mode()
+def MeanAudioInference(
+    prompt='',
+    negative_prompt='',
+    model_path='',
+    encoder_name='t5_clap',
+    variant='meanaudio_mf',
+    duration=10,
+    cfg_strength=4.5,
+    num_steps=1,
+    output='./output',
+    seed=42,
+    full_precision=False,
+    use_rope=True,
+    text_c_dim=512,
+    use_meanflow=False
+):
+    '''
+    prompt (str): 
+        The text description guiding the audio generation (e.g., "a dog is barking").
+    negative_prompt (str): 
+        A text description for sounds that should be avoided in the generated audio.
+    model_path (str): 
+        Path to the model weights file. If empty, it defaults to ./weights/{variant}.pth.
+    encoder_name (str): 
+        Specifies the text encoder to use (default: 't5_clap').
+    variant (str): 
+        Specifies the model variant to load (default: 'meanaudio_mf'). Must be a key in all_model_cfg.
+    duration (int): 
+        The desired duration of the generated audio in seconds (default: 10).
+    cfg_strength (float): 
+        Classifier-Free Guidance strength. Ignored if use_meanflow is True or variant is 'meanaudio_mf' (default: 4.5).
+    num_steps (int): 
+        Number of steps for the generation process (default: 1).
+    output (str): 
+        Directory path where the generated audio file will be saved (default: './output').
+    seed (int): 
+        Random seed for generation reproducibility (default: 42).
+    full_precision (bool): 
+        If True, uses torch.float32 precision; otherwise, uses torch.bfloat16 (default: False).
+    use_rope (bool): 
+        Whether to use Rotary Position Embedding in the model (default: True).
+    text_c_dim (int): 
+        Dimension of the text context vector (default: 512).
+    use_meanflow (bool): 
+        If True, uses the MeanFlow generation method; otherwise, uses FlowMatching. If variant is 'meanaudio_mf', this is automatically set to True (default: False).    
+    '''
+    setup_eval_logging()
+    output_dir = Path(output).expanduser()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    dtype = torch.float32 if full_precision else torch.bfloat16
+    if duration <= 0 or num_steps <= 0:
+        raise ValueError("Duration and number of steps must be positive.")
+    if variant not in all_model_cfg:
+         raise ValueError(f"Unknown model variant: {variant}. Available: {list(all_model_cfg.keys())}")
+    if not model_path or model_path == '':
+        model_path = Path(f'./weights/{variant}.pth')
+    else:
+        model_path = Path(model_path)
+    if not model_path.exists():
+        if str(model_path) == f'./weights/{variant}.pth':
+            log.info(f'Model not found at {model_path}')
+            log.info('Downloading models to "./weights/"...')
+            try:
+                weights_dir = Path('./weights')
+                weights_dir.mkdir(exist_ok=True)
+                snapshot_download(repo_id="junxiliu/Meanaudio", local_dir="./weights",allow_patterns=["*.pt", "*.pth"] )
+                raise NotImplementedError("Model download functionality needs to be implemented")
+            except Exception as e:
+                log.error(f"Failed to download model: {e}")
+                raise FileNotFoundError(f"Model file not found and download failed: {model_path}")
+        else:
+            raise FileNotFoundError(f"Model file not found: {model_path}")
+    
+    model = all_model_cfg[variant]
+    seq_cfg = model.seq_cfg
+    seq_cfg.duration = duration
+    
+    net = get_mean_audio(model.model_name, use_rope=use_rope, text_c_dim=text_c_dim)
+    net = net.to(device, dtype).eval()
+    net.load_weights(torch.load(model_path, map_location=device, weights_only=True))
+    net.update_seq_lengths(seq_cfg.latent_seq_len)
+    
+    if variant=='meanaudio_mf':
+        use_meanflow=True
+    if use_meanflow:
+        generation_func = MeanFlow(steps=num_steps)
+        cfg_strength=0
+    else:
+        generation_func = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+    
+    feature_utils = FeaturesUtils(
+        tod_vae_ckpt=model.vae_path,
+        enable_conditions=True,
+        encoder_name=encoder_name,
+        mode=model.mode,
+        bigvgan_vocoder_ckpt=model.bigvgan_16k_path,
+        need_vae_encoder=False
+    )
+    feature_utils = feature_utils.to(device, dtype).eval()
+    
+    rng = torch.Generator(device=device)
+    rng.manual_seed(seed)
+    
+    generate_fn = generate_mf if use_meanflow else generate_fm
+    kwargs = {
+        'negative_text': [negative_prompt],
+        'feature_utils': feature_utils,
+        'net': net,
+        'rng': rng,
+        'cfg_strength': cfg_strength
+    }
+    
+    if use_meanflow:
+        kwargs['mf'] = generation_func
+    else:
+        kwargs['fm'] = generation_func
+    
+    audios = generate_fn([prompt], **kwargs)
+    audio = audios.float().cpu()[0]
+    safe_filename = prompt.replace(' ', '_').replace('/', '_').replace('.', '')
+    save_path = output_dir / f'{safe_filename}--numsteps{num_steps}--seed{seed}.wav'
+    torchaudio.save(save_path, audio, seq_cfg.sampling_rate)
+    log.info(f'Audio saved to {save_path}')
+    log.info('Memory usage: %.2f GB', torch.cuda.max_memory_allocated() / (2**30))
+    return save_path
+
+if __name__ == '__main__':
+    MeanAudioInference('a dog is barking')

app.py

@@ -0,0 +1,421 @@
+
+import warnings
+
+warnings.filterwarnings("ignore", category=FutureWarning)
+import logging
+from argparse import ArgumentParser
+from pathlib import Path
+import torch
+import torchaudio
+import gradio as gr
+from meanaudio.eval_utils import (
+    ModelConfig,
+    all_model_cfg,
+    generate_mf,
+    generate_fm,
+    setup_eval_logging,
+)
+from meanaudio.model.flow_matching import FlowMatching
+from meanaudio.model.mean_flow import MeanFlow
+from meanaudio.model.networks import MeanAudio, get_mean_audio
+from meanaudio.model.utils.features_utils import FeaturesUtils
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+import gc
+from datetime import datetime
+
+log = logging.getLogger()
+
+device = "cpu"
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    log.warning("CUDA/MPS are not available, running on CPU")
+setup_eval_logging()
+
+OUTPUT_DIR = Path("./output/gradio")
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+current_model_state = {
+    "net": None,
+    "feature_utils": None,
+    "seq_cfg": None,
+    "args": None,
+}
+
+
+def load_model_if_needed(
+    variant, model_path, encoder_name, use_rope, text_c_dim, full_precision
+):
+    global current_model_state
+    dtype = torch.float32 if full_precision else torch.bfloat16
+    needs_reload = (
+        current_model_state["args"] is None
+        or current_model_state["args"].variant != variant
+        or current_model_state["args"].model_path != model_path
+        or current_model_state["args"].encoder_name != encoder_name
+        or current_model_state["args"].use_rope != use_rope
+        or current_model_state["args"].text_c_dim != text_c_dim
+        or current_model_state["args"].full_precision != full_precision
+    )
+    if needs_reload:
+        try:
+            if variant not in all_model_cfg:
+                raise ValueError(f"Unknown model variant: {variant}")
+            model: ModelConfig = all_model_cfg[variant]
+            seq_cfg = model.seq_cfg
+
+            class MockArgs:
+                pass
+
+            mock_args = MockArgs()
+            mock_args.variant = variant
+            mock_args.model_path = model_path
+            mock_args.encoder_name = encoder_name
+            mock_args.use_rope = use_rope
+            mock_args.text_c_dim = text_c_dim
+            mock_args.full_precision = full_precision
+
+            net: MeanAudio = (
+                get_mean_audio(
+                    model.model_name,
+                    use_rope=mock_args.use_rope,
+                    text_c_dim=mock_args.text_c_dim,
+                )
+                .to(device, dtype)
+                .eval()
+            )
+            net.load_weights(
+                torch.load(
+                    mock_args.model_path, map_location=device, weights_only=True
+                )
+            )
+            log.info(f"Loaded weights from {mock_args.model_path}")
+
+            feature_utils = FeaturesUtils(
+                tod_vae_ckpt=model.vae_path,
+                enable_conditions=True,
+                encoder_name=mock_args.encoder_name,
+                mode=model.mode,
+                bigvgan_vocoder_ckpt=model.bigvgan_16k_path,
+                need_vae_encoder=False,
+            )
+            feature_utils = feature_utils.to(device, dtype).eval()
+
+            current_model_state["net"] = net
+            current_model_state["feature_utils"] = feature_utils
+            current_model_state["seq_cfg"] = seq_cfg
+            current_model_state["args"] = mock_args
+            log.info(f"Model '{variant}' loaded successfully.")
+            return True
+        except Exception as e:
+            log.error(f"Error loading model: {e}")
+
+            current_model_state = {
+                "net": None,
+                "feature_utils": None,
+                "seq_cfg": None,
+                "args": None,
+            }
+            raise e
+    else:
+        log.info(f"Model '{variant}' already loaded with current settings.")
+        return False
+
+
+@torch.inference_mode()
+def generate_audio_gradio(
+    prompt,
+    negative_prompt,
+    duration,
+    cfg_strength,
+    num_steps,
+    seed,
+    variant,
+    full_precision,
+):
+    global current_model_state
+    use_meanflow = variant == "meanaudio_mf"
+
+    model_path = (
+        "./weights/meanaudio_mf.pth"
+        if use_meanflow
+        else "./weights/fluxaudio_fm.pth"
+    )
+    encoder_name = "t5_clap"
+    use_rope = True
+    text_c_dim = 512
+
+    try:
+        load_model_if_needed(
+            variant, model_path, encoder_name, use_rope, text_c_dim, full_precision
+        )
+    except Exception as e:
+        return f"Error loading model: {str(e)}", None
+
+    if current_model_state["net"] is None:
+        return "Error: Model could not be loaded.", None
+    net = current_model_state["net"]
+    feature_utils = current_model_state["feature_utils"]
+    seq_cfg = current_model_state["seq_cfg"]
+
+    args = current_model_state["args"]
+    dtype = torch.float32 if args.full_precision else torch.bfloat16
+
+    try:
+        seq_cfg.duration = duration
+        net.update_seq_lengths(seq_cfg.latent_seq_len)
+
+        rng = torch.Generator(device=device)
+        if seed >= 0:
+            rng.manual_seed(seed)
+        else:
+            rng.seed()
+
+        if use_meanflow:
+            sampler = MeanFlow(steps=num_steps)
+            log.info("Using MeanFlow for generation.")
+            generation_func = generate_mf
+            sampler_arg_name = "mf"
+            cfg_strength = 3
+        else:
+            sampler = FlowMatching(
+                min_sigma=0, inference_mode="euler", num_steps=num_steps
+            )
+            log.info("Using FlowMatching for generation.")
+            generation_func = generate_fm
+            sampler_arg_name = "fm"
+
+        prompts = [prompt]
+
+        audios = generation_func(
+            prompts,
+            negative_text=[negative_prompt],
+            feature_utils=feature_utils,
+            net=net,
+            rng=rng,
+            cfg_strength=cfg_strength,
+            **{sampler_arg_name: sampler},
+        )
+        audio = audios.float().cpu()[0]
+
+        safe_prompt = (
+            "".join(c for c in prompt if c.isalnum() or c in (" ", "_"))
+            .rstrip()
+            .replace(" ", "_")[:50]
+        )
+        current_time_string = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+        filename = f"{safe_prompt}_{current_time_string}.flac"
+        save_path = OUTPUT_DIR / filename
+        torchaudio.save(str(save_path), audio, seq_cfg.sampling_rate)
+        log.info(f"Audio saved to {save_path}")
+
+        gc.collect()
+
+        return (
+            f"Generated audio for prompt: '{prompt}' using {'MeanFlow' if use_meanflow else 'FlowMatching'}",
+            str(save_path),
+        )
+    except Exception as e:
+        gc.collect()
+        log.error(f"Generation error: {e}")
+        return f"Error during generation: {str(e)}", None
+
+
+theme = gr.themes.Soft(
+    primary_hue="blue",
+    secondary_hue="slate",
+    neutral_hue="slate",
+    text_size="sm",
+    spacing_size="sm",
+).set(
+    background_fill_primary="*neutral_50",
+    background_fill_secondary="*background_fill_primary",
+    block_background_fill="*background_fill_primary",
+    block_border_width="0px",
+    panel_background_fill="*neutral_50",
+    panel_border_width="0px",
+    input_background_fill="*neutral_100",
+    input_border_color="*neutral_200",
+    button_primary_background_fill="*primary_300",
+    button_primary_background_fill_hover="*primary_400",
+    button_secondary_background_fill="*neutral_200",
+    button_secondary_background_fill_hover="*neutral_300",
+)
+
+custom_css = """
+#main-header {
+    text-align: center;
+    margin-top: 5px;
+    margin-bottom: 10px;
+    color: var(--neutral-600);
+    font-weight: 600;
+}
+#model-settings-header, #generation-settings-header {
+    color: var(--neutral-600);
+    margin-top: 8px;
+    margin-bottom: 8px;
+    font-weight: 500;
+    font-size: 1.1em;
+}
+.setting-section {
+    padding: 10px 12px;
+    border-radius: 6px;
+    background-color: var(--neutral-50);
+    margin-bottom: 10px;
+    border: 1px solid var(--neutral-100);
+}
+hr {
+    border: none;
+    height: 1px;
+    background-color: var(--neutral-200);
+    margin: 8px 0;
+}
+#generate-btn {
+    width: 100%;
+    max-width: 250px;
+    margin: 10px auto;
+    display: block;
+    padding: 10px 15px;
+    font-size: 16px;
+    border-radius: 5px;
+}
+#status-box {
+    min-height: 50px;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    padding: 8px;
+    border-radius: 5px;
+    border: 1px solid var(--neutral-200);
+    color: var(--neutral-700);
+}
+#audio-output {
+    height: 100px;
+    border-radius: 5px;
+    border: 1px solid var(--neutral-200);
+}
+.gradio-dropdown label, .gradio-checkbox label, .gradio-number label, .gradio-textbox label {
+    font-weight: 500;
+    color: var(--neutral-700);
+    font-size: 0.9em;
+}
+.gradio-row {
+   gap: 8px;
+}
+.gradio-block {
+   margin-bottom: 8px;
+}
+.setting-section .gradio-block {
+    margin-bottom: 6px;
+}
+::-webkit-scrollbar {
+  width: 8px;
+  height: 8px;
+}
+::-webkit-scrollbar-track {
+  background: var(--neutral-100);
+  border-radius: 4px;
+}
+::-webkit-scrollbar-thumb {
+  background: var(--neutral-300);
+  border-radius: 4px;
+}
+::-webkit-scrollbar-thumb:hover {
+  background: var(--neutral-400);
+}
+* {
+  scrollbar-width: thin;
+  scrollbar-color: var(--neutral-300) var(--neutral-100);
+}
+"""
+
+with gr.Blocks(title="MeanAudio Generator", theme=theme, css=custom_css) as demo:
+    gr.Markdown("# MeanAudio Text-to-Audio Generator", elem_id="main-header")
+
+    gr.Markdown("### Model and Generation Settings", elem_id="model-settings-header")
+    with gr.Column(elem_classes="setting-section"):
+        with gr.Row():
+            available_variants = (
+                list(all_model_cfg.keys()) if all_model_cfg else []
+            )
+            default_variant = (
+                "small_16k_mf"
+                if "small_16k_mf" in available_variants
+                else available_variants[0] if available_variants else ""
+            )
+            variant = gr.Dropdown(
+                label="Model Variant",
+                choices=available_variants,
+                value=default_variant,
+                interactive=True,
+                scale=3,
+            )
+            full_precision = gr.Checkbox(
+                label="Full Precision (float32)", value=True, scale=1
+            )
+
+    gr.Markdown("### Audio Generation", elem_id="generation-settings-header")
+    with gr.Column(elem_classes="setting-section"):
+        with gr.Row():
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Describe the sound you want to generate...",
+                scale=1,
+            )
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="Describe sounds you want to avoid...",
+                value="",
+                scale=1,
+            )
+        with gr.Row():
+            duration = gr.Number(
+                label="Duration (sec)", value=10.0, minimum=0.1, scale=1
+            )
+            cfg_strength = gr.Number(
+                label="CFG (Meanflow forced to 3)", value=3, minimum=0.0, scale=1
+            )
+        with gr.Row():
+            seed = gr.Number(
+                label="Seed (-1 for random)", value=42, precision=0, scale=1
+            )
+            num_steps = gr.Number(
+                label="Number of Steps",
+                value=1,
+                precision=0,
+                minimum=1,
+                scale=1,
+            )
+
+    generate_button = gr.Button("Generate", variant="primary", elem_id="generate-btn")
+    generate_output_text = gr.Textbox(
+        label="Result Status", interactive=False, elem_id="status-box"
+    )
+    audio_output = gr.Audio(
+        label="Generated Audio", type="filepath", elem_id="audio-output"
+    )
+    generate_button.click(
+        fn=generate_audio_gradio,
+        inputs=[
+            prompt,
+            negative_prompt,
+            duration,
+            cfg_strength,
+            num_steps,
+            seed,
+            variant,
+            full_precision,
+        ],
+        outputs=[generate_output_text, audio_output],
+    )
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--port", type=int, default=7861)
+    args = parser.parse_args()
+    demo.launch(server_port=args.port, allowed_paths=[OUTPUT_DIR.resolve()])
+

easyinfer.py

@@ -0,0 +1,3 @@
+from MeanAudio import MeanAudioInference
+audio_path=MeanAudioInference('a dog is barking')
+print(audio_path)

requirements.txt

@@ -0,0 +1,27 @@
+torch>=2.5.1
+huggingface_hub>=0.26
+cython
+gitpython>=3.1
+tensorboard>=2.11
+numpy>=1.21,<2.1
+Pillow>=9.5
+opencv-python>=4.8
+scipy>=1.7
+tqdm>=4.66.1
+gradio>=3.34
+einops>=0.6
+hydra-core>=1.3.2
+requests
+torchdiffeq>=0.2.5
+librosa>=0.8.1
+nitrous-ema
+hydra_colorlog
+tensordict>=0.6.1
+colorlog
+open_clip_torch>=2.29.0
+av>=14.0.1
+timm>=1.0.12
+python-dotenv
+transformers
+debugpy
+laion-clap