change model

app.py

@@ -1,32 +1,28 @@
 from fastapi import FastAPI
-from transformers import CsmForConditionalGeneration, AutoProcessor
+from transformers import AutoProcessor
+from peft import PeftModel
+from transformers import CsmForConditionalGeneration
 import torch
+import soundfile as sf
 import aiofiles
 import uvicorn
 import os
-import soundfile as sf
 
 app = FastAPI()
 
-# Load model and processor once
-MODEL_NAME = "Ellbendls/csm-1b-indonesian-fine-tuned"
+BASE_MODEL = "deep-io/csm-1b"
+FINETUNE_MODEL = "Ellbendls/csm-1b-indonesian-fine-tuned"
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = CsmForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
-processor = AutoProcessor.from_pretrained(MODEL_NAME)
 
-def generate_timestamps(script: str, wpm: int = 150):
-    sentences = [s.strip() for s in script.replace("\n", " ").split(".") if s.strip()]
-    timestamps, current = [], 0.0
-    spw = 60.0 / wpm
-    for s in sentences:
-        dur = round(len(s.split()) * spw, 2)
-        timestamps.append({"sentence": s, "start": round(current,2), "end": round(current+dur,2)})
-        current += dur
-    return timestamps
+# Load base model
+model = CsmForConditionalGeneration.from_pretrained(BASE_MODEL, torch_dtype=torch.float16).to(device)
+
+# Apply LoRA adapter
+model = PeftModel.from_pretrained(model, FINETUNE_MODEL)
 
-@app.get("/")
-async def root():
-    return {"message": "API is running!"}
+# Load processor
+processor = AutoProcessor.from_pretrained(BASE_MODEL)
 
 @app.post("/tts")
 async def tts_api(payload: dict):
@@ -34,20 +30,19 @@ async def tts_api(payload: dict):
     if not text:
         return {"error": "Text is required"}
 
-    file_path = "output.wav"
-    input_text = f"[0]{text}"
-    inputs = processor(input_text, return_tensors="pt").to(device)
-    outputs = model.generate(**inputs, max_new_tokens=200, output_audio=True)
-    audio = outputs[0].detach().cpu().numpy()
+    inputs = processor(f"[0]{text}", return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=200, output_audio=True)
 
+    audio = outputs[0].detach().cpu().numpy()
+    file_path = "output.wav"
     sf.write(file_path, audio, 24000)
 
     async with aiofiles.open(file_path, "rb") as f:
         audio_data = await f.read()
 
     return {
-        "script": text,
-        "timestamps": generate_timestamps(text),
+        "text": text,
         "file_url": f"/download/{os.path.basename(file_path)}",
         "size": len(audio_data)
     }

requirements.txt

@@ -2,7 +2,8 @@ fastapi
 uvicorn[standard]
 aiofiles
 pydantic
-transformers
+transformers>=4.41.0
 torch
 soundfile
-sentencepiece
+sentencepiece
+peft