fix issue

app.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, File, UploadFile, Response
-from transformers import AutoModelForSpeechSeq2Seq, AutoTokenizer
+from transformers import ParlerTTSForConditionalGeneration, AutoTokenizer
 from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
 from llama_cpp import Llama
 import torch
@@ -10,74 +10,59 @@ from pydantic import BaseModel
 
 app = FastAPI()
 
-# Load TTS model
-TTS_MODEL_PATH = "./models/tts_model"
-if os.path.exists(TTS_MODEL_PATH):
-    tts_model = AutoModelForSpeechSeq2Seq.from_pretrained(TTS_MODEL_PATH)
-    tts_tokenizer = AutoTokenizer.from_pretrained(TTS_MODEL_PATH)
+# Load models
+if os.path.exists("./models/tts_model"):
+    tts_model = ParlerTTSForConditionalGeneration.from_pretrained("./models/tts_model")
+    tts_tokenizer = AutoTokenizer.from_pretrained("./models/tts_model")
 else:
-    tts_model = AutoModelForSpeechSeq2Seq.from_pretrained("suno/bark")  # Replace with an actual TTS model
-    tts_tokenizer = AutoTokenizer.from_pretrained("suno/bark")
+    tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-mini-v1")
+    tts_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-mini-v1")
 
-# Load SST model
-SST_MODEL_PATH = "./models/sst_model"
-if os.path.exists(SST_MODEL_PATH):
-    sst_model = Wav2Vec2ForCTC.from_pretrained(SST_MODEL_PATH)
-    sst_processor = Wav2Vec2Processor.from_pretrained(SST_MODEL_PATH)
+# SST and LLM loading remains unchanged
+if os.path.exists("./models/sst_model"):
+    sst_model = Wav2Vec2ForCTC.from_pretrained("./models/sst_model")
+    sst_processor = Wav2Vec2Processor.from_pretrained("./models/sst_model")
 else:
     sst_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
-    sst_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+    sst_processor = Wav2Vec2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
 
-# Load Llama.cpp model
-LLM_MODEL_PATH = "./models/llama.gguf"
-if os.path.exists(LLM_MODEL_PATH):
-    llm = Llama(model_path=LLM_MODEL_PATH)  # Corrected usage
+if os.path.exists("./models/llama.gguf"):
+    llm = Llama("./models/llama.gguf")
 else:
     raise FileNotFoundError("Please upload llama.gguf to models/ directory")
 
-# Request models
+# Request models and endpoints remain unchanged
 class TTSRequest(BaseModel):
     text: str
 
 class LLMRequest(BaseModel):
     prompt: str
 
-# API Endpoints
 @app.post("/tts")
 async def tts_endpoint(request: TTSRequest):
     text = request.text
     inputs = tts_tokenizer(text, return_tensors="pt")
-    
     with torch.no_grad():
-        output_ids = tts_model.generate(**inputs)
-    
-    # Convert model output to speech (assuming Bark-like model)
-    audio = output_ids.squeeze().cpu().numpy()
-    
+        audio = tts_model.generate(**inputs)
+    audio = audio.squeeze().cpu().numpy()
     buffer = io.BytesIO()
-    sf.write(buffer, audio, samplerate=22050, format="WAV")  # Ensure correct sample rate
+    sf.write(buffer, audio, 22050, format="WAV")
     buffer.seek(0)
-    
     return Response(content=buffer.getvalue(), media_type="audio/wav")
 
 @app.post("/sst")
 async def sst_endpoint(file: UploadFile = File(...)):
     audio_bytes = await file.read()
     audio, sr = sf.read(io.BytesIO(audio_bytes))
-    
     inputs = sst_processor(audio, sampling_rate=sr, return_tensors="pt")
-    
     with torch.no_grad():
-        logits = sst_model(inputs.input_values).logits  # Ensure .logits is detached
-        predicted_ids = torch.argmax(logits, dim=-1)
-    
+        logits = sst_model(inputs.input_values).logits
+    predicted_ids = torch.argmax(logits, dim=-1)
     transcription = sst_processor.batch_decode(predicted_ids)[0]
-    
     return {"text": transcription}
 
 @app.post("/llm")
 async def llm_endpoint(request: LLMRequest):
     prompt = request.prompt
-    output = llm(prompt, max_tokens=50)  # Corrected llama.cpp API call
-    
-    return {"text": output["choices"][0]["text"] if "choices" in output else output["content"]}
+    output = llm(prompt, max_tokens=50)
+    return {"text": output["choices"][0]["text"]}