Update app.py

app.py

@@ -14,6 +14,7 @@ from transformers import (
 )
 import boto3
 import uvicorn
+import asyncio
 from transformers import pipeline
 import json
 from huggingface_hub import login
@@ -77,7 +78,7 @@ class S3ModelLoader:
          return f"s3://{self.bucket_name}/" \
                f"{model_name.replace('/', '-')}"
 
-    def load_model_and_tokenizer(self, model_name):
+    async def load_model_and_tokenizer(self, model_name):
        if model_name in model_cache:
           return model_cache[model_name]
        
@@ -147,7 +148,7 @@ class S3ModelLoader:
 model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 
 @app.post("/generate")
-def generate(request: GenerateRequest):
+async def generate(request: GenerateRequest):
     try:
         model_name = request.model_name
         input_text = request.input_text
@@ -162,7 +163,7 @@ def generate(request: GenerateRequest):
         do_sample = request.do_sample
         stop_sequences = request.stop_sequences
 
-        model, tokenizer = model_loader.load_model_and_tokenizer(model_name)
+        model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
 
@@ -185,7 +186,7 @@ def generate(request: GenerateRequest):
                         media_type="text/plain"
                     )
             else:
-                result = generate_text(model, tokenizer, input_text,
+                result = await generate_text(model, tokenizer, input_text,
                                      generation_config, stop_sequences,
                                      device)
                 return JSONResponse({"text": result, "is_end": True})
@@ -212,7 +213,7 @@ class StopOnSequences(StoppingCriteria):
                  return True
         return False
 
-def stream_text(model, tokenizer, input_text,
+async def stream_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
                         device):
     
@@ -228,7 +229,7 @@ def stream_text(model, tokenizer, input_text,
     
     while True:
         
-        outputs = model.generate(
+        outputs = await asyncio.to_thread(model.generate,
                                           **encoded_input,
                                            do_sample=generation_config.do_sample,
                                             max_new_tokens=generation_config.max_new_tokens,
@@ -270,7 +271,7 @@ def stream_text(model, tokenizer, input_text,
         output_text = ""
         
 
-def generate_text(model, tokenizer, input_text,
+async def generate_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
                         device):
     encoded_input = tokenizer(
@@ -281,7 +282,7 @@ def generate_text(model, tokenizer, input_text,
     stop_criteria = StopOnSequences(stop_sequences, tokenizer)
     stopping_criteria = StoppingCriteriaList([stop_criteria])
     
-    outputs = model.generate(
+    outputs = await asyncio.to_thread(model.generate,
                                       **encoded_input,
                                        do_sample=generation_config.do_sample,
                                         max_new_tokens=generation_config.max_new_tokens,
@@ -302,8 +303,9 @@ def generate_text(model, tokenizer, input_text,
     
     return generated_text
 
+
 @app.post("/generate-image")
-def generate_image(request: GenerateRequest):
+async def generate_image(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -331,7 +333,7 @@ def generate_image(request: GenerateRequest):
 
 
 @app.post("/generate-text-to-speech")
-def generate_text_to_speech(request: GenerateRequest):
+async def generate_text_to_speech(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -362,7 +364,7 @@ def generate_text_to_speech(request: GenerateRequest):
 
 
 @app.post("/generate-video")
-def generate_video(request: GenerateRequest):
+async def generate_video(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"