update

app.py

@@ -1,528 +0,0 @@
-import warnings
-import spaces
-warnings.filterwarnings("ignore", category=FutureWarning)
-import logging
-from argparse import ArgumentParser
-from pathlib import Path
-import torch
-import torchaudio
-import gradio as gr
-from transformers import AutoModel
-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
-from huggingface_hub import snapshot_download
-
-log = logging.getLogger()
-device = "cpu"
-if torch.cuda.is_available():
-    device = "cuda"
-setup_eval_logging()
-OUTPUT_DIR = Path("./output/gradio")
-OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-
-# --- 预下载依赖 ---
-snapshot_download(repo_id="google/flan-t5-large")
-#snapshot_download(repo_id="google-bert/bert-base-uncased")
-a=AutoModel.from_pretrained('bert-base-uncased')
-b=AutoModel.from_pretrained('roberta-base')
-#snapshot_download(repo_id="FacebookAI/roberta-base")
-snapshot_download(repo_id="junxiliu/Meanaudio", local_dir="./weights",allow_patterns=["*.pt", "*.pth"] )
-
-# --- 模型状态存储 (修改为字典存储多个模型) ---
-# Stores states for multiple models, keyed by variant name
-current_model_states = {
-    # Example structure:
-    # "meanaudio_mf": { "net": ..., "feature_utils": ..., "seq_cfg": ..., "args": ... },
-    # "fluxaudio_fm": { "net": ..., "feature_utils": ..., "seq_cfg": ..., "args": ... },
-}
-
-# --- 模型加载逻辑 (修改以支持多模型) ---
-def load_model_if_needed(
-    variant, model_path, encoder_name, use_rope, text_c_dim
-):
-    global current_model_states
-    dtype = torch.float32 
-
-    # Check if this specific variant with these args is already loaded
-    existing_state = current_model_states.get(variant)
-    needs_reload = (
-        existing_state is None
-        or existing_state["args"].variant != variant
-        or existing_state["args"].model_path != model_path
-        or existing_state["args"].encoder_name != encoder_name
-        or existing_state["args"].use_rope != use_rope
-        or existing_state["args"].text_c_dim != text_c_dim
-        
-    )
-
-    if needs_reload:
-        log.info(f"Loading/reloading model '{variant}'.")
-        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
-
-        # Create mock args object
-        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
-        
-
-        # Load network
-        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}")
-
-        # Load feature utils
-        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()
-
-        # Store the loaded model state
-        current_model_states[variant] = {
-            "net": net,
-            "feature_utils": feature_utils,
-            "seq_cfg": seq_cfg,
-            "args": mock_args,
-        }
-        log.info(f"Model '{variant}' loaded successfully.")
-        # Return the loaded components for immediate use if needed
-        return net, feature_utils, seq_cfg, mock_args
-
-    else:
-        log.info(f"Model '{variant}' already loaded with current settings. Skipping reload.")
-        # Return the existing components
-        return existing_state["net"], existing_state["feature_utils"], existing_state["seq_cfg"], existing_state["args"]
-
-# --- 初始化函数，用于预加载所有默认模型 ---
-def initialize_all_default_models():
-    """Function to load all default model configurations at startup."""
-    log.info("Initializing default models...")
-    default_models = ['meanaudio_mf', 'fluxaudio_fm'] # List of default variants to load
-    common_params = {
-        "encoder_name": "t5_clap",
-        "use_rope": True,
-        "text_c_dim": 512,
-         # Match the default value in the UI or desired startup precision
-    }
-
-    for variant in default_models:
-        model_path = f"./weights/{variant}.pth"
-        # This will load the model if it's not already loaded or if params differ
-        try:
-            load_model_if_needed(
-                variant, model_path, **common_params
-            )
-            log.info(f"Default model '{variant}' initialized successfully.")
-        except Exception as e:
-            log.error(f"Failed to initialize default model '{variant}': {e}")
-            # Depending on requirements, decide if failure to load one model should stop the app
-            # For now, we'll just log and continue trying others.
-initialize_all_default_models()
-# --- 带有 GPU 装饰器的生成函数 (修改为使用预加载模型) ---
-@spaces.GPU(duration=8) # This decorator now primarily handles moving the *already loaded* model/functionality to the GPU instance for execution
-@torch.inference_mode()
-def generate_audio_gradio(
-    prompt,
-    negative_prompt,
-    duration,
-    cfg_strength,
-    num_steps,
-    seed,
-    variant, # This determines which pre-loaded model to use
-    
-):
-    global current_model_states
-    # Determine model parameters based on input (mainly for path)
-    model_path = f"./weights/{variant}.pth" # Determine path based on variant
-    encoder_name = "t5_clap"
-    use_rope = True
-    text_c_dim = 512
-
-    # --- 关键修改点：获取已加载的模型组件 ---
-    # Fetch the pre-loaded model components based on the selected variant
-    model_state = current_model_states.get(variant)
-    if model_state is None:
-        # This case should ideally not happen if initialization was successful,
-        # but handle it gracefully in case of unexpected state.
-        error_msg = f"Error: Model '{variant}' is not available. It may not have been loaded correctly during startup."
-        log.error(error_msg)
-        return error_msg, None
-
-    # Use the pre-loaded components
-    net = model_state["net"]
-    feature_utils = model_state["feature_utils"]
-    seq_cfg = model_state["seq_cfg"]
-    # Use the args stored with the model for consistency (e.g., dtype)
-    args = model_state["args"]
-    # Ensure dtype consistency based on the loaded model's args or UI input
-    dtype = torch.float32  # Use UI input for dtype consistency
-
-    # --- Rest of your generation logic ---
-    # Update sequence length based on duration
-    temp_seq_cfg = type(seq_cfg)(**seq_cfg.__dict__) # Create a temporary copy to modify duration
-    temp_seq_cfg.duration = duration
-    # Update network sequence lengths
-    net.update_seq_lengths(temp_seq_cfg.latent_seq_len)
-
-    # Setup random number generator
-    rng = torch.Generator(device=device)
-    if seed >= 0:
-        rng.manual_seed(seed)
-    else:
-        rng.seed()
-
-    # Select sampler and generation function based on variant
-    use_meanflow = variant == "meanaudio_mf"
-    if use_meanflow:
-        sampler = MeanFlow(steps=num_steps)
-        log.info("Using MeanFlow for generation.")
-        generation_func = generate_mf
-        sampler_arg_name = "mf"
-        # Note: cfg_strength is forced to 3 for MeanFlow inside the function
-        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"
-
-    # Perform generation
-    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] # Get the first generated audio and move to CPU
-
-    # Save the generated audio
-    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, temp_seq_cfg.sampling_rate) # Use temp_seq_cfg for correct SR
-    log.info(f"Audio saved to {save_path}")
-
-    # Cleanup
-    gc.collect()
-    # torch.cuda.empty_cache() # Optional: if using CUDA and want to be more aggressive
-
-    return (
-        f"Generated audio for prompt: '{prompt}' using {'MeanFlow' if use_meanflow else 'FlowMatching'}",
-        str(save_path),
-    )
-
-
-# --- Gradio UI 和事件 ---
-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-headertitle {
-    text-align: center;
-    margin-top: 15px;
-    margin-bottom: 10px;
-    color: var(--neutral-600);
-    font-weight: 600;
-}
-#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);
-}
-#project-badges {
-    text-align: center; /* 内容居中 */
-    margin-top: 30px;
-    margin-bottom: 20px; /* 与下方内容的间距 */
-}
-
-/* 针对徽章容器 div */
-#project-badges #badge-container {
-    display: flex;
-    gap: 10px;
-    align-items: center;
-    justify-content: center;
-    flex-wrap: wrap;
-}
-
-/* 针对徽章图片本身 */
-#project-badges img { /* 使用 #project-badges img 可以更精确地只影响这个区域的图片 */
-    border-radius: 5px; /* 圆角 */
-    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1); /* 阴影 */
-    height: 20px; /* 统一高度 */
-    transition: transform 0.1s ease, box-shadow 0.1s ease; /* 添加悬停效果的过渡 */
-}
-
-/* 可选：添加悬停效果 */
-#project-badges a:hover img {
-    transform: translateY(-2px); /* 向上轻微移动 */
-    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.15); /* 加强阴影 */
-}
-#audio-output {
-    height: 200px;
-    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:Fast and Faithful Text-to-Audio Generation with Mean Flows", elem_id="main-header")
-    #gr.Markdown("### Model and Generation Settings", elem_id="model-settings-header")
-    #[![arXiv](https://img.shields.io/badge/arXiv-2412.21037-b31b1b.svg?logo=arXiv)](https://arxiv.org/abs/2412.21037)
-    project_badges_markdown = '''
-<div style="display: flex; gap: 10px; align-items: center; justify-content: center; flex-wrap: wrap; margin-bottom: 20px;">
-  <a href="https://huggingface.co/junxiliu/MeanAudio">
-    <img src="https://img.shields.io/badge/Model-HuggingFace-violet?logo=huggingface" alt="Hugging Face Model">
-  </a>
-  <a href="https://huggingface.co/spaces/chenxie95/MeanAudio">
-    <img src="https://img.shields.io/badge/Space-HuggingFace-8A2BE2?logo=huggingface" alt="Hugging Face Space">
-  </a>
-  <a href="https://meanaudio.github.io/">
-    <img src="https://img.shields.io/badge/Project-Page-brightred?style=flat" alt="Project Page">
-  </a>
-  <a href="https://github.com/xiquan-li/MeanAudio">
-    <img src="https://img.shields.io/badge/Code-GitHub-black?logo=github" alt="GitHub">
-  </a>
-</div>
-'''
-    # 使用 gr.Markdown 渲染徽章行
-    gr.Markdown(project_badges_markdown, elem_id="project-badges")
-    with gr.Column(elem_classes="setting-section"):
-        with gr.Row():
-            available_variants = (
-                list(all_model_cfg.keys()) if all_model_cfg else []
-            )
-            default_variant = (
-                'meanaudio_mf'
-            )
-            variant = gr.Dropdown(
-                label="Model Variant",
-                choices=available_variants,
-                value=default_variant,
-                interactive=True,
-                scale=3,
-            )
-    #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,
-            
-        ],
-        outputs=[generate_output_text, audio_output],
-    )
-    audio_examples = [
-        # [prompt, negative_prompt, duration, cfg_strength, num_steps, seed, variant]
-        ["A speech and gunfire followed by a gun being loaded", "", 10.0, 3, 1, 42, "meanaudio_mf"],
-        ["Typing on a keyboard", "", 10.0, 3, 1, 42, "meanaudio_mf"],
-        ["A man speaks followed by a popping noise and laughter", "", 10.0, 3, 2, 42, "meanaudio_mf"],
-        ["Some humming followed by a toilet flushing", "", 10.0, 3, 2, 42, "meanaudio_mf"],
-        ["Rain falling on a hard surface as thunder roars in the distance", "", 10.0, 3, 5, 42, "meanaudio_mf"],
-        ["Food sizzling and oil popping", "", 10.0, 3, 25, 42, "meanaudio_mf"],
-        ["Pots and dishes clanking as a man talks followed by liquid pouring into a container", "", 8.0, 3, 2, 42, "meanaudio_mf"],
-        ["A few seconds of silence then a rasping sound against wood", "", 12.0, 3, 2, 42, "meanaudio_mf"],
-        ["A man speaks as he gives a speech and then the crowd cheers", "", 10.0, 3, 25, 42, "fluxaudio_fm"],
-        ["A goat bleating repeatedly", "", 10.0, 3, 50, 123, "fluxaudio_fm"],
-        ["Tires squealing followed by an engine revving", "", 12.0, 4, 25, 456, "fluxaudio_fm"],
-        ["Hammer slowly hitting the wooden table", "", 10.0, 3.5, 25, 42, "fluxaudio_fm"], 
-        ["Dog barking excitedly and man shouting as race car engine roars past", "", 10.0, 3, 1, 42, "meanaudio_mf"],
-        ["A dog barking and a cat mewing and a racing car passes by", "", 12.0, 3, 5, -1, "meanaudio_mf"],
-        ["Whistling with birds chirping", "", 10.0, 4, 50, 42, "fluxaudio_fm"],
-
-    ]
-    gr.Examples(
-        examples=audio_examples,
-        inputs=[prompt, negative_prompt, duration, cfg_strength, num_steps, seed, variant], # 必须与示例数据顺序一致
-        outputs=[generate_output_text, audio_output], # 指定输出组件
-        fn=generate_audio_gradio, # 指定处理函数
-        # cache_examples=True, # 可选：缓存示例结果（需要 Hugging Face Space 支持或特定设置）
-        examples_per_page=5, # 可选：每页显示的示例数量
-        label="Example Prompts", # 可选：示例组的标签
-    )
-    # --- 关键修改点：使用 load 事件进行初始化，加载所有默认模型 ---
-    # This ensures all default models are loaded when the app starts
-    #demo.load(fn=initialize_all_default_models, inputs=None, outputs=None)
-
-if __name__ == "__main__":
-    demo.launch()