Update main.py

main.py

@@ -1,59 +1,38 @@
-import spaces
 import os
+import gc
 import torch
 import soundfile as sf
 import logging
 import gradio as gr
 import librosa
 import numpy as np
+import spaces
+
 from datetime import datetime
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from lmdeploy import pipeline, GenerationConfig, TurbomindEngineConfig
 from ncodec.codec import TTSCodec
 
-# ----------------- Logging -----------------
+# ---------------- Logging ----------------
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
-# ----------------- Globals -----------------
-MODEL_PIPE = None
-TOKENIZER = None
+# ---------------- Globals ----------------
+GPU_PIPE = None
 CODEC = None
+MODEL_ID = "rahul7star/mir-TTS"
 
-# ----------------- Model Initialization (CPU ONLY) -----------------
-def initialize_model():
-    global MODEL_PIPE, TOKENIZER, CODEC
-
-    if MODEL_PIPE is not None:
-        return MODEL_PIPE
-
-    logging.info("Loading MiraTTS model on CPU...")
-
-    model_name = "rahul7star/mir-TTS"
-
-    TOKENIZER = AutoTokenizer.from_pretrained(model_name)
-
-    MODEL_PIPE = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float32,   # CPU safe
-        device_map=None
-    )
-
-    MODEL_PIPE.eval()
-    MODEL_PIPE.to("cpu")  # 🔒 CPU only
-
-    CODEC = TTSCodec()
 
-    logging.info("Model loaded successfully on CPU")
-    return MODEL_PIPE
+# ---------------- CPU Init (SAFE) ----------------
+def initialize_cpu():
+    global CODEC
+    if CODEC is None:
+        logging.info("Initializing CPU components")
+        CODEC = TTSCodec()
 
 
-# 🔹 Load model ONCE at startup (CPU safe)
-MODEL_PIPE = initialize_model()
-
-
-# ----------------- Audio Utilities -----------------
+# ---------------- Audio Utils ----------------
 def validate_audio_input(audio_path):
     if not audio_path or not os.path.exists(audio_path):
         raise ValueError("Audio file not found")
@@ -69,77 +48,95 @@ def validate_audio_input(audio_path):
 
     audio = audio / np.max(np.abs(audio))
 
-    temp_path = f"/tmp/processed_{os.path.basename(audio_path)}"
-    sf.write(temp_path, audio, sr)
+    tmp_path = f"/tmp/processed_{os.path.basename(audio_path)}"
+    sf.write(tmp_path, audio, sr)
 
-    return temp_path
+    return tmp_path
 
 
-# ----------------- TTS Generation (GPU ONLY) -----------------
+# ---------------- GPU TTS ----------------
 @spaces.GPU()
 def generate_speech(text, audio_path):
-    global MODEL_PIPE
+    global GPU_PIPE, CODEC
 
     if not text or not text.strip():
         raise ValueError("Text input is empty")
 
-    # 🔥 Move model to GPU only here
-    if not next(MODEL_PIPE.parameters()).is_cuda:
-        logging.info("Moving model to GPU for generation")
-        MODEL_PIPE.to("cuda")
+    initialize_cpu()
+
+    # 🔥 Load GPU pipeline lazily (CORRECT)
+    if GPU_PIPE is None:
+        logging.info("Loading MiraTTS pipeline on GPU")
+
+        backend_config = TurbomindEngineConfig(
+            tp=1,
+            device="cuda",
+            dtype="bfloat16",
+            enable_prefix_caching=False,
+            cache_max_entry_count=0.1,
+        )
+
+        GPU_PIPE = pipeline(
+            MODEL_ID,
+            backend_config=backend_config
+        )
 
     processed_audio = validate_audio_input(audio_path)
     context_tokens = CODEC.encode(processed_audio)
 
     prompt = CODEC.format_prompt(text, context_tokens, None)
 
-    inputs = TOKENIZER(prompt, return_tensors="pt").to("cuda")
+    gen_cfg = GenerationConfig(
+        top_p=0.95,
+        top_k=50,
+        temperature=0.8,
+        max_new_tokens=1024,
+        repetition_penalty=1.2,
+        do_sample=True,
+    )
 
-    with torch.no_grad():
-        outputs = MODEL_PIPE.generate(
-            **inputs,
-            max_new_tokens=1024,
-            top_p=0.95,
-            top_k=50,
-            temperature=0.8,
-            repetition_penalty=1.2,
-        )
+    response = GPU_PIPE(
+        [prompt],
+        gen_config=gen_cfg,
+        do_preprocess=False
+    )
 
-    generated_text = TOKENIZER.decode(outputs[0], skip_special_tokens=True)
-    audio = CODEC.decode(generated_text, context_tokens)
+    audio = CODEC.decode(response[0].text, context_tokens)
 
     if torch.is_tensor(audio):
-        audio = audio.cpu().numpy()
+        audio = audio.float().cpu().numpy()  # force float32
 
-    # 🧹 Cleanup GPU memory
-    del inputs, outputs
+    # 🧹 Cleanup
+    os.remove(processed_audio)
+    gc.collect()
     torch.cuda.empty_cache()
 
     return audio, 48000
 
 
-# ----------------- Gradio Interface -----------------
+# ---------------- Gradio ----------------
 def voice_clone_interface(text, upload_audio, record_audio):
     try:
-        audio_path = upload_audio if upload_audio else record_audio
+        audio_path = upload_audio or record_audio
         if not audio_path:
-            return None, "Please upload or record reference audio."
+            return None, "Upload or record reference audio"
 
         audio, sr = generate_speech(text, audio_path)
 
         os.makedirs("outputs", exist_ok=True)
-        out_path = f"outputs/mira_{datetime.now().strftime('%Y%m%d_%H%M%S')}.wav"
+        out_path = f"outputs/mira_{datetime.now():%Y%m%d_%H%M%S}.wav"
         sf.write(out_path, audio, sr)
 
         return out_path, "✅ Generation successful"
 
     except Exception as e:
-        return None, f"❌ Error: {str(e)}"
+        logging.error(e)
+        return None, f"❌ {str(e)}"
 
 
 def build_interface():
     with gr.Blocks(title="MiraTTS Voice Cloning") as demo:
-        gr.Markdown("# 🎤 MiraTTS Voice Cloning")
+        gr.Markdown("# 🎤 MiraTTS – Voice Cloning")
 
         with gr.Row():
             with gr.Column():
@@ -162,7 +159,8 @@ def build_interface():
     return demo
 
 
-# ----------------- Main -----------------
+# ---------------- Main ----------------
 if __name__ == "__main__":
+    initialize_cpu()
     demo = build_interface()
     demo.launch(server_name="0.0.0.0", server_port=7860)