Update app.py

app.py

@@ -85,24 +85,14 @@ def load_models():
     # Load tokenizer
     mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
     
-    # Load codec model
-    model_config = OmegaConf.load(CODEC_CONFIG_PATH)
-    codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(DEVICE)
-    parameter_dict = torch.load(CODEC_CKPT_PATH, map_location='cpu')
-    codec_model.load_state_dict(parameter_dict['codec_model'])
-    codec_model.eval()
-    
-    # Initialize codec tools
-    codectool = CodecManipulator("xcodec", 0, 1)
-    
     # Precompute token IDs
     start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
     end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
     
-    return model, mmtokenizer, codec_model, codectool, start_of_segment, end_of_segment
+    return model, mmtokenizer, start_of_segment, end_of_segment
 
 # Preload all models and components
-model, mmtokenizer, codec_model, codectool, start_of_segment, end_of_segment = load_models()
+model, mmtokenizer, start_of_segment, end_of_segment = load_models()
 
 # Audio processing cache
 resampler_cache = {}
@@ -134,6 +124,16 @@ def generate_music(
     keep_intermediate=False,
     rescale=False,
 ):
+    # Load codec model
+    model_config = OmegaConf.load(CODEC_CONFIG_PATH)
+    codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(DEVICE)
+    parameter_dict = torch.load(CODEC_CKPT_PATH, map_location='cpu')
+    codec_model.load_state_dict(parameter_dict['codec_model'])
+    codec_model.eval()
+    
+    # Initialize codec tools
+    codectool = CodecManipulator("xcodec", 0, 1)
+    
     # Create output directories once
     os.makedirs(output_dir, exist_ok=True)
     stage1_output_dir = os.path.join(output_dir, "stage1")