Update app.py

app.py

@@ -42,18 +42,25 @@ MODEL_CONFIG = {
 # -------------------------------
 
 class ModelManager:
+    def __init__(self):
+        self.llama_pipelines = {}
+        self.musicgen_model = None
+        self.tts_models = {}
+        self.processor = None  # Add processor cache
+
     def get_llama_pipeline(self, model_id, token):
         if model_id not in self.llama_pipelines:
             tokenizer = AutoTokenizer.from_pretrained(
                 model_id,
                 token=token,
-                legacy=False  # Critical for tokenizers 0.19.x compatibility
+                legacy=False
             )
             model = AutoModelForCausalLM.from_pretrained(
                 model_id,
                 token=token,
                 torch_dtype=torch.float16,
-                device_map="auto"
+                device_map="auto",
+                low_cpu_mem_usage=True
             )
             self.llama_pipelines[model_id] = pipeline(
                 "text-generation",
@@ -62,19 +69,15 @@ class ModelManager:
                 device_map="auto"
             )
         return self.llama_pipelines[model_id]
-        
+
     def get_musicgen_model(self):
         if not self.musicgen_model:
             self.musicgen_model = MusicgenForConditionalGeneration.from_pretrained(
                 MODEL_CONFIG["musicgen_model"]
             )
+            self.processor = AutoProcessor.from_pretrained(MODEL_CONFIG["musicgen_model"])
             self.musicgen_model.to("cuda" if torch.cuda.is_available() else "cpu")
-        return self.musicgen_model
-
-    def get_tts_model(self, model_name):
-        if model_name not in self.tts_models:
-            self.tts_models[model_name] = TTS(model_name)
-        return self.tts_models[model_name]
+        return self.musicgen_model, self.processor
 
 model_manager = ModelManager()