Update app.py

app.py

@@ -3,6 +3,8 @@ from fastapi.responses import FileResponse
 import torch
 import torchaudio
 import os
+from pydantic import BaseModel
+from typing import List, Optional
 
 
 # MODEL_DIR = "my_model"
@@ -50,6 +52,32 @@ model.load_checkpoint(
 )
 model.to(device)
 
+# --------- Define your models ----------
+class BGM(BaseModel):
+    file: str
+    bgm_volume: float
+
+class Sentence(BaseModel):
+    sentence_id: int
+    speaker: str
+    text: str
+    prosody_ref: str
+
+class Scene(BaseModel):
+    scene_id: int
+    ambiance: str
+    bgm: BGM
+    sentences: List[Sentence]
+
+class CastMember(BaseModel):
+    name: str
+    gender: str
+    voice_ref: str
+
+class StoryInput(BaseModel):
+    cast: List[CastMember]
+    scenes: List[Scene]
+
 def tts_arabic(text: str, audio_file: str) -> str:
     gpt_cond_latent, speaker_embedding = model.get_conditioning_latents(audio_path=[audio_file])
     out = model.inference(