Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import torch
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel, field_validator
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList
 import boto3
@@ -55,16 +55,16 @@ class GenerateRequest(BaseModel):
 
     @field_validator("max_new_tokens")
     def max_new_tokens_must_be_within_limit(cls, v):
-        if v > 4:
-            raise ValueError("max_new_tokens cannot be greater than 4.")
+        if v > 500:
+            raise ValueError("max_new_tokens cannot be greater than 500.")
         return v
 
 class S3ModelLoader:
-    def.__init__(self, bucket_name, s3_client):
+    def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
         self.s3_client = s3_client
 
-    def._get_s3_uri(self, model_name):
+    def _get_s3_uri(self, model_name):
         return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
 
     async def load_model_and_tokenizer(self, model_name):
@@ -187,7 +187,10 @@ async def generate_text_to_speech(request: GenerateRequest):
         audio = audio_generator(validated_body.input_text)[0]
 
         audio_byte_arr = BytesIO()
-        audio.save(audio_byte_arr)
+        # It is expected that the audio is saved as wav.
+        # Saving like this will not always work. Please check how your
+        # audio_generator model is working.
+        audio_generator.save_audio(audio_byte_arr, audio)
         audio_byte_arr.seek(0)
 
         return StreamingResponse(audio_byte_arr, media_type="audio/wav")
@@ -204,7 +207,10 @@ async def generate_video(request: GenerateRequest):
         video = video_generator(validated_body.input_text)[0]
 
         video_byte_arr = BytesIO()
-        video.save(video_byte_arr)
+        # Same as above. Please check how your video model is returning the
+        # videos and save them accordingly.
+        # It is expected that the video is saved as MP4
+        video_generator.save_video(video_byte_arr, video)
         video_byte_arr.seek(0)
 
         return StreamingResponse(video_byte_arr, media_type="video/mp4")