update

app.py

@@ -38,18 +38,48 @@ OUTPUT_DIR = Path("./output/gradio")
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 NUM_SAMPLE = 1
 
-# snapshot_download(repo_id="google/flan-t5-large")
-# a = AutoModel.from_pretrained('bert-base-uncased')
-# b = AutoModel.from_pretrained('roberta-base')
+# Global model cache to avoid reloading
+MODEL_CACHE = {}
+FEATURE_UTILS_CACHE = {}
 
-# 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()
+def ensure_models_downloaded():
+    for variant, model_cfg in all_model_cfg.items():
+        if not model_cfg.model_path.exists():
+            log.info(f'Model {variant} not found, downloading...')
+            snapshot_download(repo_id="AndreasXi/MeanAudio", local_dir="./weights", allow_patterns=["*.pt", "*.pth"])
+            break
 
-# laion_clap_model.load_ckpt(_clap_ckpt_path, verbose=False)
+def load_model_if_needed(variant: str):
+    if variant in MODEL_CACHE:
+        return MODEL_CACHE[variant], FEATURE_UTILS_CACHE[variant]
+    
+    log.info(f"Loading model {variant} for the first time...")
+    model_cfg = all_model_cfg[variant]
+    
+    net = get_mean_audio(model_cfg.model_name, use_rope=True, text_c_dim=512)
+    net = net.to(device, torch.bfloat16).eval()
+    net.load_weights(torch.load(model_cfg.model_path, map_location=device, weights_only=True))
+    
+    feature_utils = FeaturesUtils(
+        tod_vae_ckpt=model_cfg.vae_path,
+        enable_conditions=True,
+        encoder_name="t5_clap", 
+        mode=model_cfg.mode,
+        bigvgan_vocoder_ckpt=model_cfg.bigvgan_16k_path,
+        need_vae_encoder=False
+    )
+    feature_utils = feature_utils.to(device, torch.bfloat16).eval()
+    
+    MODEL_CACHE[variant] = net
+    FEATURE_UTILS_CACHE[variant] = feature_utils
+    
+    log.info(f"Model {variant} loaded and cached successfully")
+    return net, feature_utils
+
+ensure_models_downloaded()
 
 
-@spaces.GPU(duration=10)
+@spaces.GPU(duration=60)
 @torch.inference_mode()
 def generate_audio_gradio(
     prompt,
@@ -66,29 +96,14 @@ def generate_audio_gradio(
     if variant not in all_model_cfg:
         raise ValueError(f"Unknown model variant: {variant}. Available: {list(all_model_cfg.keys())}")
 
-    model_path = all_model_cfg[variant].model_path  # by default, this will use meanaudio_s_full.pth or fluxaudio_s_full.pth
-    if not model_path.exists():
-        log.info(f'Model not found at {model_path}')
-        log.info('Downloading models to "./weights/"...')
-        snapshot_download(repo_id="AndreasXi/MeanAudio", local_dir="./weights",allow_patterns=["*.pt", "*.pth"] )
-
+    net, feature_utils = load_model_if_needed(variant)
+    
     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)
 
-    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()
-
 
     if variant == 'meanaudio_s_ac' or variant == 'meanaudio_s_full':
         use_meanflow=True
@@ -141,7 +156,8 @@ def generate_audio_gradio(
     torchaudio.save(str(save_path), audio, seq_cfg.sampling_rate)
     log.info(f"Audio saved to {save_path}")
 
-    gc.collect()
+    if device == "cuda":
+        torch.cuda.empty_cache()
 
     return (
         f"Generated audio for prompt: '{prompt}' using {'MeanFlow' if use_meanflow else 'FlowMatching'}",