modified: app.py

app.py

@@ -82,17 +82,10 @@ def load_models():
         attn_implementation="flash_attention_2",
     ).to(DEVICE).eval()
     
-    # Load tokenizer
-    mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-    
-    # Precompute token IDs
-    start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
-    end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
-    
-    return model, mmtokenizer, start_of_segment, end_of_segment
+    return model
 
 # Preload all models and components
-model, mmtokenizer, start_of_segment, end_of_segment = load_models()
+model = load_models()
 
 # Audio processing cache
 resampler_cache = {}
@@ -124,6 +117,13 @@ def generate_music(
     keep_intermediate=False,
     rescale=False,
 ):
+    # Load tokenizer
+    mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
+    
+    # Precompute token IDs
+    start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
+    end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
+    
     # Load codec model
     model_config = OmegaConf.load(CODEC_CONFIG_PATH)
     codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(DEVICE)