Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import logging
 import threading
 import boto3
-from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, StoppingCriteriaList
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, StoppingCriteriaList, pipeline
 from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel, field_validator
 from huggingface_hub import hf_hub_download
@@ -77,7 +77,7 @@ class S3ModelLoader:
     def download_model_from_huggingface(self, model_name):
         try:
             logging.info(f"Downloading model {model_name} from Hugging Face...")
-            model_dir = hf_hub_download(model_name, token=HUGGINGFACE_HUB_TOKEN, filename=None)
+            model_dir = hf_hub_download(model_name, token=HUGGINGFACE_HUB_TOKEN)
             model_files = os.listdir(model_dir)
             for model_file in model_files:
                 s3_path = f"lilmeaty_garca/{model_name}/{model_file}"
@@ -85,6 +85,7 @@ class S3ModelLoader:
             logging.info(f"Model {model_name} saved to S3 successfully.")
         except Exception as e:
             logging.error(f"Error downloading model {model_name} from Hugging Face: {e}")
+            raise HTTPException(status_code=500, detail=f"Error downloading model from Hugging Face: {e}")
 
     def download_all_models_in_background(self):
         models_url = "https://huggingface.co/api/models"
@@ -104,6 +105,17 @@ class S3ModelLoader:
     def run_in_background(self):
         threading.Thread(target=self.download_all_models_in_background, daemon=True).start()
 
+    def load_model_and_tokenizer(self, model_name):
+        try:
+            model_uri = self._download_from_s3(model_name)
+            model = AutoModelForCausalLM.from_pretrained(model_uri)
+            tokenizer = AutoTokenizer.from_pretrained(model_uri)
+            logging.info(f"Model {model_name} loaded successfully from {model_uri}.")
+            return model, tokenizer
+        except Exception as e:
+            logging.error(f"Error loading model {model_name}: {e}")
+            raise HTTPException(status_code=500, detail=f"Error loading model {model_name}: {e}")
+
 @app.on_event("startup")
 async def startup_event():
     model_loader.run_in_background()
@@ -174,29 +186,21 @@ async def generate(request: Request, body: GenerateRequest):
             generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device)
             audio = generator(validated_body.input_text)
             audio_bytesio = BytesIO()
-            sf.write(audio_bytesio, audio["sampling_rate"], np.int16(audio["audio"]))
-            audio_bytes = audio_bytesio.getvalue()
-            return Response(content=audio_bytes, media_type="audio/wav")
+            sf.write(audio_bytesio, audio["samples"], audio["rate"], format="WAV")
+            audio_bytesio.seek(0)
+            return StreamingResponse(audio_bytesio, media_type="audio/wav")
 
         elif validated_body.task_type == "text-to-video":
-            try:
-                generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device)
-                video = generator(validated_body.input_text)
-                return Response(content=video, media_type="video/mp4")
-            except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Error generating video: {str(e)}")
-
+            return {"error": "Text-to-video task type is not yet supported."}
         else:
-            raise HTTPException(status_code=400, detail="Invalid task type.")
+            raise HTTPException(status_code=400, detail="Invalid task type")
 
     except Exception as e:
-        logging.error(f"Error processing request: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        logging.error(f"Error during generation: {e}")
+        raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")
+
+import uvicorn
+
+if __name__ == "__main__":
+    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
 
-def download_model_from_s3_or_hf(model_name):
-    try:
-        model_dir = model_loader._download_from_s3(model_name)
-        return model_dir
-    except Exception:
-        model_loader.download_model_from_huggingface(model_name)
-        return model_loader._download_from_s3(model_name)