Update app.py

app.py

@@ -2,17 +2,15 @@ import os
 import logging
 import threading
 import boto3
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, StoppingCriteriaList, AutoConfig
 from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel, field_validator
-from io import BytesIO
 from huggingface_hub import hf_hub_download
 import requests
+import time
 import asyncio
-import soundfile as sf
-import numpy as np
 from fastapi.responses import StreamingResponse, Response
+import torch
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
 
@@ -67,30 +65,22 @@ class S3ModelLoader:
             if "Contents" not in model_files:
                 raise FileNotFoundError(f"Model files not found in S3 for {model_name}")
 
-            local_dir = f"/tmp/{model_name.replace('/', '-')}"
-            os.makedirs(local_dir, exist_ok=True)
-
-            for obj in model_files["Contents"]:
-                file_key = obj["Key"]
-                if file_key.endswith('/'):
-                    continue
-
-                local_file_path = os.path.join(local_dir, os.path.basename(file_key))
-                self.s3_client.download_file(self.bucket_name, file_key, local_file_path)
-
-            return local_dir
+            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
         except Exception as e:
             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:
-            model_dir = await asyncio.to_thread(self._download_from_s3, model_name)
-            config = AutoConfig.from_pretrained(model_dir)
-            tokenizer = AutoTokenizer.from_pretrained(model_dir, config=config)
-            model = AutoModelForCausalLM.from_pretrained(model_dir, config=config)
+            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 from S3 successfully.")
+            logging.info(f"Model {model_name} loaded successfully from S3.")
             return model, tokenizer
         except Exception as e:
             logging.exception(f"Error loading model: {e}")
@@ -128,6 +118,7 @@ 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,13 +195,15 @@ async def generate(request: Request, body: GenerateRequest):
                 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 in text-to-video generation: {e}")
+                raise HTTPException(status_code=500, detail=f"Error generating video: {str(e)}")
 
         else:
-            raise HTTPException(status_code=400, detail="Unsupported task type")
+            raise HTTPException(status_code=400, detail="Invalid task type.")
 
-    except HTTPException as e:
-        raise e
     except Exception as e:
-        logging.exception(f"An unexpected error occurred: {e}")
-        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
+        logging.error(f"Error processing request: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)