update new model versions and test

app.py

@@ -16,6 +16,7 @@ from meanaudio.eval_utils import (
     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
@@ -25,117 +26,28 @@ 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)
 NUM_SAMPLE=7
-snapshot_download(repo_id="google/flan-t5-large")
-a=AutoModel.from_pretrained('bert-base-uncased')
-b=AutoModel.from_pretrained('roberta-base')
-snapshot_download(repo_id="junxiliu/Meanaudio", local_dir="./weights",allow_patterns=["*.pt", "*.pth"] )
-_clap_ckpt_path='./weights/music_speech_audioset_epoch_15_esc_89.98.pt'
-laion_clap_model = laion_clap.CLAP_Module(enable_fusion=False,
-                                              amodel='HTSAT-base').cuda().eval()
-laion_clap_model.load_ckpt(_clap_ckpt_path, verbose=False)
-current_model_states = {
-
-}
-
-def load_model_if_needed(
-    variant, model_path, encoder_name, use_rope, text_c_dim
-):
-    global current_model_states
-    dtype = torch.float32
-    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
-
-        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
-
-        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_states[variant] = {
-            "net": net,
-            "feature_utils": feature_utils,
-            "seq_cfg": seq_cfg,
-            "args": mock_args,
-        }
-        log.info(f"Model '{variant}' loaded successfully.")
-
-        return net, feature_utils, seq_cfg, mock_args
-    else:
-        log.info(f"Model '{variant}' already loaded with current settings. Skipping reload.")
-
-        return existing_state["net"], existing_state["feature_utils"], existing_state["seq_cfg"], existing_state["args"]
 
-def initialize_all_default_models():
-    log.info("Initializing default models...")
-    default_models = ['meanaudio_mf', 'fluxaudio_fm']
-    common_params = {
-        "encoder_name": "t5_clap",
-        "use_rope": True,
-        "text_c_dim": 512,
+snapshot_download(repo_id="google/flan-t5-large")
+a = AutoModel.from_pretrained('bert-base-uncased')
+b = AutoModel.from_pretrained('roberta-base')
 
-    }
-    for variant in default_models:
-        model_path = f"./weights/{variant}.pth"
+snapshot_download(repo_id="AndreasXi/MeanAudio", local_dir="./weights",allow_patterns=["*.pt", "*.pth"] )
+_clap_ckpt_path='./weights/music_speech_audioset_epoch_15_esc_89.98.pt'
+laion_clap_model = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-base').cuda().eval()
 
-        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}")
+laion_clap_model.load_ckpt(_clap_ckpt_path, verbose=False)
 
-initialize_all_default_models()
 
 @spaces.GPU(duration=10)
 @torch.inference_mode()
@@ -148,44 +60,42 @@ def generate_audio_gradio(
     seed,
     variant,
 ):
-    global current_model_states
-
-    model_path = f"./weights/{variant}.pth"
-    encoder_name = "t5_clap"
-    use_rope = True
-    text_c_dim = 512
-
-    model_state = current_model_states.get(variant)
-    if model_state is None:
-        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
-
-    net = model_state["net"]
-    feature_utils = model_state["feature_utils"]
-    seq_cfg = model_state["seq_cfg"]
-
-    args = model_state["args"]
     dtype = torch.float32
+    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())}")
 
-    temp_seq_cfg = type(seq_cfg)(**seq_cfg.__dict__)
-    temp_seq_cfg.duration = duration
+    model_path = all_model_cfg[variant].model_path  # by default, this will use meanaudio_s_full.pth or fluxaudio_s_full.pth
+    model = all_model_cfg[variant]
+    seq_cfg = model.seq_cfg
+    seq_cfg.duration = duration
+    
+    net = get_mean_audio(model.model_name, use_rope=True, text_c_dim=512)
+    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)
 
-    net.update_seq_lengths(temp_seq_cfg.latent_seq_len)
+    feature_utils = FeaturesUtils(tod_vae_ckpt=model.vae_path,
+                                  enable_conditions=True,
+                                  encoder_name="t5_clap", 
+                                  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)
-    if seed >= 0:
-        rng.manual_seed(seed)
-    else:
-        rng.seed()
 
-    use_meanflow = variant == "meanaudio_mf"
+    if variant == 'meanaudio_s_ac' or variant == 'meanaudio_s_full':
+        use_meanflow=True
+    elif variant == 'fluxaudio_s_full':
+        use_meanflow=False
+
     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
+        cfg_strength = 0
     else:
         sampler = FlowMatching(
             min_sigma=0, inference_mode="euler", num_steps=num_steps
@@ -193,6 +103,10 @@ def generate_audio_gradio(
         log.info("Using FlowMatching for generation.")
         generation_func = generate_fm
         sampler_arg_name = "fm"
+
+    rng = torch.Generator(device=device)
+    # rng.manual_seed(seed)
+
     audios = generation_func(
         [prompt]*NUM_SAMPLE,
         negative_text=[negative_prompt]*NUM_SAMPLE,
@@ -205,11 +119,11 @@ def generate_audio_gradio(
     text_embed = laion_clap_model.get_text_embedding(prompt, use_tensor=True).squeeze()
     audio_embed = laion_clap_model.get_audio_embedding_from_data(audios[:,0,:].float().cpu(), use_tensor=True).squeeze()
     scores = torch.cosine_similarity(text_embed.expand(audio_embed.shape[0], -1),
-                                audio_embed,
-                                dim=-1)
+                                     audio_embed,
+                                     dim=-1)
     log.info(scores)
     log.info(torch.argmax(scores).item())
-    audio=audios[torch.argmax(scores).item()].float().cpu()
+    audio = audios[torch.argmax(scores).item()].float().cpu()
     safe_prompt = (
         "".join(c for c in prompt if c.isalnum() or c in (" ", "_"))
         .rstrip()
@@ -400,7 +314,6 @@ with gr.Blocks(title="MeanAudio Generator", theme=theme, css=custom_css) as demo
                 interactive=True,
                 scale=3,
             )
-
     with gr.Column(elem_classes="setting-section"):
         with gr.Row():
             prompt = gr.Textbox(

meanaudio/eval_utils.py

@@ -43,20 +43,26 @@ class ModelConfig:
             download_model_if_needed(self.bigvgan_16k_path)
 
 
-fluxaudio_fm = ModelConfig(model_name='fluxaudio_fm', 
-                           model_path=Path('./weights/fluxaudio_fm.pth'),
+fluxaudio_s_full = ModelConfig(model_name='fluxaudio_s_full', 
+                           model_path=Path('./weights/fluxaudio_s_full.pth'),  # will be specified later 
                            vae_path=Path('./weights/v1-16.pth'),
                            bigvgan_16k_path=Path('./weights/best_netG.pt'),
                            mode='16k')
-meanaudio_mf = ModelConfig(model_name='meanaudio_mf', 
-                           model_path=Path('./weights/meanaudio_mf.pth'),
+meanaudio_s_full = ModelConfig(model_name='meanaudio_s_full', 
+                           model_path=Path('./weights/meanaudio_s_full.pth'),  # will be specified later 
+                           vae_path=Path('./weights/v1-16.pth'),
+                           bigvgan_16k_path=Path('./weights/best_netG.pt'),
+                           mode='16k')
+meanaudio_s_ac = ModelConfig(model_name='meanaudio_s_ac', 
+                           model_path=Path('./weights/meanaudio_s_ac.pth'),  # will be specified later 
                            vae_path=Path('./weights/v1-16.pth'),
                            bigvgan_16k_path=Path('./weights/best_netG.pt'),
                            mode='16k')
 
 all_model_cfg: dict[str, ModelConfig] = {
-    'fluxaudio_fm': fluxaudio_fm, 
-    'meanaudio_mf': meanaudio_mf, 
+    'fluxaudio_s_full': fluxaudio_s_full, 
+    'meanaudio_s_full': meanaudio_s_full, 
+    'meanaudio_s_ac': meanaudio_s_ac, 
 }
 
 

meanaudio/model/networks.py

@@ -577,7 +577,7 @@ class MeanAudio(nn.Module):
         return self._latent_seq_len
     
 
-def fluxaudio_fm(**kwargs) -> FluxAudio: 
+def fluxaudio_s(**kwargs) -> FluxAudio: 
     num_heads = 7
     return FluxAudio(latent_dim=20,
                      text_dim=1024,
@@ -587,7 +587,7 @@ def fluxaudio_fm(**kwargs) -> FluxAudio:
                      num_heads=num_heads,
                      latent_seq_len=312,  # for 10s audio
                      **kwargs)
-def meanaudio_mf(**kwargs) -> MeanAudio: 
+def meanaudio_s(**kwargs) -> MeanAudio: 
     num_heads = 7
     return MeanAudio(latent_dim=20,
                      text_dim=1024,
@@ -600,10 +600,10 @@ def meanaudio_mf(**kwargs) -> MeanAudio:
 
 
 def get_mean_audio(name: str, **kwargs) -> MeanAudio:
-    if name == 'meanaudio_mf':
-        return meanaudio_mf(**kwargs)
-    if name == 'fluxaudio_fm': 
-        return fluxaudio_fm(**kwargs)
+    if name == 'meanaudio_s':
+        return meanaudio_s(**kwargs)
+    if name == 'fluxaudio_s': 
+        return fluxaudio_s(**kwargs)
 
     raise ValueError(f'Unknown model name: {name}')
 
@@ -620,7 +620,7 @@ if __name__ == '__main__':
         ]
     )
 
-    network: MeanAudio = get_mean_audio('meanaudio_mf', 
+    network: MeanAudio = get_mean_audio('meanaudio_s', 
                                         use_rope=False, 
                                         text_c_dim=512)