Update app.py

app.py

@@ -11,6 +11,9 @@ import time
 import asyncio
 from fastapi.responses import StreamingResponse, Response
 import torch
+from io import BytesIO
+import numpy as np
+import soundfile as sf
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
 
@@ -64,7 +67,6 @@ class S3ModelLoader:
             model_files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=f"lilmeaty_garca/{model_name}")
             if "Contents" not in model_files:
                 raise FileNotFoundError(f"Model files not found in S3 for {model_name}")
-
             s3_model_path = f"s3://{self.bucket_name}/lilmeaty_garca/{model_name.replace('/', '-')}"
             logging.info(f"Model {model_name} found on S3 at {s3_model_path}")
             return s3_model_path
@@ -72,25 +74,14 @@ class S3ModelLoader:
             logging.error(f"Error downloading from S3: {e}")
             raise HTTPException(status_code=500, detail=f"Error downloading model from S3: {e}")
 
-    async def load_model_and_tokenizer(self, model_name):
-        try:
-            s3_model_path = await asyncio.to_thread(self._download_from_s3, model_name)
-            # Load from S3 directly (no local storage)
-            config = AutoConfig.from_pretrained(s3_model_path)
-            tokenizer = AutoTokenizer.from_pretrained(s3_model_path, config=config)
-            model = AutoModelForCausalLM.from_pretrained(s3_model_path, config=config)
-
-            logging.info(f"Model {model_name} loaded successfully from S3.")
-            return model, tokenizer
-        except Exception as e:
-            logging.exception(f"Error loading model: {e}")
-            raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
-
     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=model_name.split("/")[-1])
-            self.s3_client.upload_file(model_dir, self.bucket_name, f"lilmeaty_garca/{model_name}")
+            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}"
+                self.s3_client.upload_file(os.path.join(model_dir, model_file), self.bucket_name, s3_path)
             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}")
@@ -102,7 +93,6 @@ class S3ModelLoader:
             if response.status_code != 200:
                 logging.error("Error getting Hugging Face model list.")
                 raise HTTPException(status_code=500, detail="Error getting model list.")
-
             models = response.json()
             for model in models:
                 model_name = model["id"]
@@ -118,7 +108,6 @@ class S3ModelLoader:
 async def startup_event():
     model_loader.run_in_background()
 
-# Initialize S3 client with boto3
 s3_client = boto3.client('s3', aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY, region_name=AWS_REGION)
 model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 
@@ -204,6 +193,21 @@ async def generate(request: Request, body: GenerateRequest):
         logging.error(f"Error processing request: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
+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)
+
+def ensure_s3_directories(model_name):
+    try:
+        s3_path = f"lilmeaty_garca/{model_name}"
+        s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=s3_path)
+    except Exception as e:
+        logging.error(f"Error ensuring S3 directories exist for model {model_name}: {e}")
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)