modified: app.py

app.py

@@ -76,36 +76,33 @@ MAX_SEQ_LEN = 16384
 # --------------------------
 # Preload Models with KV Cache Initialization
 # --------------------------
-@spaces.GPU
-def preload_models():
-    global model, mmtokenizer, codec_model, codectool, vocal_decoder, inst_decoder
-    
-    # Text generation model with KV cache support
-    model = AutoModelForCausalLM.from_pretrained(
-        "m-a-p/YuE-s1-7B-anneal-en-cot",
-        torch_dtype=TORCH_DTYPE,
-        attn_implementation="flash_attention_2",
-        use_cache=True  # Enable KV caching
-    ).to(DEVICE).eval()
 
-    # Tokenizer and codec tools
-    mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-    codectool = CodecManipulator("xcodec", 0, 1)
-    
-    # Audio codec model
-    model_config = OmegaConf.load(MODEL_DIR/"final_ckpt/config.yaml")
-    codec_model = SoundStream(**model_config.generator.config).to(DEVICE)
-    codec_model.load_state_dict(
-        torch.load(MODEL_DIR/"final_ckpt/ckpt_00360000.pth", map_location='cpu')['codec_model']
-    )
-    codec_model.eval()
+# Text generation model with KV cache support
+model = AutoModelForCausalLM.from_pretrained(
+    "m-a-p/YuE-s1-7B-anneal-en-cot",
+    torch_dtype=TORCH_DTYPE,
+    attn_implementation="flash_attention_2",
+    use_cache=True  # Enable KV caching
+).to(DEVICE).eval()
 
-    # Vocoders
-    vocal_decoder, inst_decoder = build_codec_model(
-        MODEL_DIR/"decoders/config.yaml",
-        MODEL_DIR/"decoders/decoder_131000.pth",
-        MODEL_DIR/"decoders/decoder_151000.pth"
-    )
+# Tokenizer and codec tools
+mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
+codectool = CodecManipulator("xcodec", 0, 1)
+
+# Audio codec model
+model_config = OmegaConf.load(MODEL_DIR/"final_ckpt/config.yaml")
+codec_model = SoundStream(**model_config.generator.config).to(DEVICE)
+codec_model.load_state_dict(
+    torch.load(MODEL_DIR/"final_ckpt/ckpt_00360000.pth", map_location='cpu')['codec_model']
+)
+codec_model.eval()
+
+# Vocoders
+vocal_decoder, inst_decoder = build_codec_model(
+    MODEL_DIR/"decoders/config.yaml",
+    MODEL_DIR/"decoders/decoder_131000.pth",
+    MODEL_DIR/"decoders/decoder_151000.pth"
+)
 
 # --------------------------
 # Optimized Generation with KV Cache Management
@@ -237,9 +234,6 @@ def infer(genre, lyrics, num_segments=2, max_tokens=2000):
     with tempfile.TemporaryDirectory() as tmpdir:
         return generate_music(genre, lyrics, num_segments, max_tokens)
 
-# Initialize models at startup
-preload_models()
-
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# YuE Music Generator with KV Cache Optimization")