Update app.py

app.py

@@ -10,29 +10,23 @@ from transformers import (
     GenerationConfig,
     StoppingCriteria,
     StoppingCriteriaList,
+    pipeline
 )
-import boto3
 import uvicorn
 import asyncio
 import json
+import base64
 from huggingface_hub import login
 from botocore.exceptions import NoCredentialsError
+from functools import lru_cache
 
 
-AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
-AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
-AWS_REGION = os.getenv("AWS_REGION")
-S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
 HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
-
 if HUGGINGFACE_HUB_TOKEN:
     login(token=HUGGINGFACE_HUB_TOKEN,
           add_to_git_credential=False)
 
-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)
 
 app = FastAPI()
 
@@ -66,67 +60,42 @@ class GenerateRequest(BaseModel):
 
 model_data = {}  # Global dictionary to store model data
 
-class S3ModelLoader:
-    def __init__(self, bucket_name, s3_client):
-        self.bucket_name = bucket_name
-        self.s3_client = s3_client
-
-    def _get_s3_uri(self, model_name):
-         return f"s3://{self.bucket_name}/" \
-               f"{model_name.replace('/', '-')}"
-
-    async def load_model_and_tokenizer(self, model_name):
-       if model_name in model_data:
-          return model_data[model_name]["model"], model_data[model_name]["tokenizer"]
-       
-       s3_uri = self._get_s3_uri(model_name)
-       try:
-           
-            config = AutoConfig.from_pretrained(
-                s3_uri, local_files_only=False
-            )
-            
-            model = AutoModelForCausalLM.from_pretrained(
-                s3_uri, config=config, local_files_only=False
-            )
-            
-            tokenizer = AutoTokenizer.from_pretrained(
-                s3_uri, config=config, local_files_only=False
-            )
-
-            if tokenizer.eos_token_id is not None and \
-               tokenizer.pad_token_id is None:
-                tokenizer.pad_token_id = config.pad_token_id \
-                                        or tokenizer.eos_token_id
-            model_data[model_name] = {"model":model, "tokenizer":tokenizer}
-            return model, tokenizer
-       except (EnvironmentError, NoCredentialsError):
-            try:
-                config = AutoConfig.from_pretrained(
-                    model_name, token=HUGGINGFACE_HUB_TOKEN
-                )
-                tokenizer = AutoTokenizer.from_pretrained(
-                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
-                )
-                
-                model = AutoModelForCausalLM.from_pretrained(
-                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
-                )
+model_load_lock = asyncio.Lock() # Lock to avoid race conditions
 
+@lru_cache(maxsize=None)
+async def _load_model_and_tokenizer(model_name):
+    try:
+          config = AutoConfig.from_pretrained(
+              model_name, token=HUGGINGFACE_HUB_TOKEN
+          )
+          tokenizer = AutoTokenizer.from_pretrained(
+              model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
+          )
+          
+          model = AutoModelForCausalLM.from_pretrained(
+              model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
+          )
+
+          if tokenizer.eos_token_id is not None and \
+              tokenizer.pad_token_id is None:
+              tokenizer.pad_token_id = config.pad_token_id \
+                                      or tokenizer.eos_token_id
+          
+          return {"model":model, "tokenizer":tokenizer}
+    except Exception as e:
+        raise HTTPException(
+            status_code=500, detail=f"Error loading model: {e}"
+        )
 
-                if tokenizer.eos_token_id is not None and \
-                   tokenizer.pad_token_id is None:
-                    tokenizer.pad_token_id = config.pad_token_id \
-                                            or tokenizer.eos_token_id
-                
-                model_data[model_name] = {"model":model, "tokenizer":tokenizer}
-                return model, tokenizer
-            except Exception as e:
-                raise HTTPException(
-                    status_code=500, detail=f"Error loading model: {e}"
-                )
+async def load_model_and_tokenizer(model_name):
+    async with model_load_lock:
+      if model_name in model_data:
+          return model_data[model_name].get("model"), model_data[model_name].get("tokenizer")
+      
+      model_bundle = await _load_model_and_tokenizer(model_name)
+      model_data[model_name] = model_bundle
+      return model_bundle.get("model"), model_bundle.get("tokenizer")
 
-model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 
 @app.post("/generate")
 async def generate(request: GenerateRequest):
@@ -144,8 +113,8 @@ async def generate(request: GenerateRequest):
         do_sample = request.do_sample
         stop_sequences = request.stop_sequences
 
-        model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model, tokenizer = await load_model_and_tokenizer(model_name)
+        device = "cpu" # Force CPU
         model.to(device)
 
         if "text-to-text" == task_type:
@@ -231,18 +200,17 @@ async def stream_text(model, tokenizer, input_text,
 
         if len(new_text) == 0:
             if not stop_criteria(outputs.sequences, None):
-                for text in output_text.split():
-                    yield json.dumps({"text": text, "is_end": False}) + "\n"
-                yield json.dumps({"text": "", "is_end": True}) + "\n"
+                 yield {"text": output_text, "is_end": False}
+                
+                 yield {"text": "", "is_end": True}
             break
         
         output_text += new_text
         
-        for text in new_text.split():
-              yield json.dumps({"text": text, "is_end": False}) + "\n"
+        yield {"text": new_text, "is_end": False}
               
         if stop_criteria(outputs.sequences, None):
-            yield json.dumps({"text": "", "is_end": True}) + "\n"
+            yield {"text": "", "is_end": True}
             break
         
         encoded_input = tokenizer(
@@ -250,8 +218,12 @@ async def stream_text(model, tokenizer, input_text,
             truncation=True
         ).to(device)
         output_text = ""
-        
 
+
+async def stream_json_responses(generator):
+    async for data in generator:
+        yield json.dumps(data) + "\n"
+        
 async def generate_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
                         device):
@@ -288,7 +260,7 @@ async def generate_text(model, tokenizer, input_text,
 async def generate_image(request: GenerateRequest):
     try:
         validated_body = request
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = "cpu" # Force CPU
         
         if validated_body.model_name not in model_data:
             config = AutoConfig.from_pretrained(
@@ -306,7 +278,7 @@ async def generate_image(request: GenerateRequest):
         
         image_data = list(image.getdata())
         
-        return json.dumps({"image_data": image_data, "is_end": True})
+        return JSONResponse({"image_data": image_data, "is_end": True})
 
     except Exception as e:
         raise HTTPException(
@@ -319,7 +291,7 @@ async def generate_image(request: GenerateRequest):
 async def generate_text_to_speech(request: GenerateRequest):
     try:
         validated_body = request
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = "cpu"  # Force CPU
 
         if validated_body.model_name not in model_data:
              config = AutoConfig.from_pretrained(
@@ -341,7 +313,7 @@ async def generate_text_to_speech(request: GenerateRequest):
         
         audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
         
-        return json.dumps({"audio": audio_base64, "is_end": True})
+        return JSONResponse({"audio": audio_base64, "is_end": True})
 
     except Exception as e:
         raise HTTPException(
@@ -354,7 +326,7 @@ async def generate_text_to_speech(request: GenerateRequest):
 async def generate_video(request: GenerateRequest):
     try:
         validated_body = request
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = "cpu"  # Force CPU
         if validated_body.model_name not in model_data:
             config = AutoConfig.from_pretrained(
                     validated_body.model_name, token=HUGGINGFACE_HUB_TOKEN
@@ -373,7 +345,7 @@ async def generate_video(request: GenerateRequest):
         
         video_base64 = base64.b64encode(video).decode('utf-8')
         
-        return json.dumps({"video": video_base64, "is_end": True})
+        return JSONResponse({"video": video_base64, "is_end": True})
 
     except Exception as e:
         raise HTTPException(
@@ -381,5 +353,27 @@ async def generate_video(request: GenerateRequest):
             detail=f"Internal server error: {str(e)}"
         )
 
+@app.on_event("startup")
+async def startup_event():
+    # Load models here
+    print("Loading models...")
+    
+    models_to_load = set()
+    
+    for env_var_key, env_var_value in os.environ.items():
+        if env_var_key.startswith("MODEL_NAME_"):
+            models_to_load.add(env_var_value)
+            
+    
+    for model_name in models_to_load:
+        try:
+            await load_model_and_tokenizer(model_name)
+            print(f"Model {model_name} loaded")
+        except Exception as e:
+             print(f"Error loading model {model_name}: {e}")
+            
+
+    print("Models loaded.")
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)