Update app.py

app.py

@@ -42,10 +42,10 @@ class GenerateRequest(BaseModel):
     input_text: str = ""
     task_type: str
     temperature: float = 1.0
-    max_new_tokens: int = 10
+    max_new_tokens: int = 3
     stream: bool = True
     top_p: float = 1.0
-    top_k: int = 50
+    top_k: int = 50    
     repetition_penalty: float = 1.0
     num_return_sequences: int = 1
     do_sample: bool = True
@@ -65,6 +65,8 @@ class GenerateRequest(BaseModel):
             raise ValueError(f"task_type must be one of: {valid_types}")
         return v
 
+model_cache = {}
+
 class S3ModelLoader:
     def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
@@ -75,8 +77,11 @@ class S3ModelLoader:
                f"{model_name.replace('/', '-')}"
 
     async def load_model_and_tokenizer(self, model_name):
-        s3_uri = self._get_s3_uri(model_name)
-        try:
+       if model_name in model_cache:
+          return model_cache[model_name]
+       
+       s3_uri = self._get_s3_uri(model_name)
+       try:
             config = AutoConfig.from_pretrained(
                 s3_uri, local_files_only=False
             )
@@ -93,9 +98,9 @@ class S3ModelLoader:
                tokenizer.pad_token_id is None:
                 tokenizer.pad_token_id = config.pad_token_id \
                                         or tokenizer.eos_token_id
-
+            model_cache[model_name] = (model, tokenizer)
             return model, tokenizer
-        except (EnvironmentError, NoCredentialsError):
+       except (EnvironmentError, NoCredentialsError):
             try:
                 config = AutoConfig.from_pretrained(
                     model_name, token=HUGGINGFACE_HUB_TOKEN
@@ -130,7 +135,7 @@ class S3ModelLoader:
                 tokenizer = AutoTokenizer.from_pretrained(
                     s3_uri, config=config, local_files_only=False
                 )
-                
+                model_cache[model_name] = (model, tokenizer)
                 return model, tokenizer
             except Exception as e:
                 raise HTTPException(
@@ -188,11 +193,10 @@ async def generate(request: GenerateRequest):
 async def stream_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
                         device):
-    max_length=10 #Define the max length to cut the text and generate another response
-    
+
     encoded_input = tokenizer(
         input_text, return_tensors="pt",
-        truncation=True, max_length=max_length
+        truncation=True
     ).to(device)
 
 
@@ -206,18 +210,6 @@ async def stream_text(model, tokenizer, input_text,
     
     output_text = ""
     while True:
-       
-        input_length = encoded_input["input_ids"].shape[1]
-        remaining_tokens = max_length - input_length
-        
-        if remaining_tokens <=0:
-           yield json.dumps({"text": "", "is_end": True}) + "\n"
-           break 
-           
-        generation_config.max_new_tokens = min(
-        remaining_tokens, generation_config.max_new_tokens
-        )
-
         outputs = model.generate(
             **encoded_input,
             do_sample=generation_config.do_sample,
@@ -242,48 +234,45 @@ async def stream_text(model, tokenizer, input_text,
 
         if stop_index != -1:
             final_output = output_text[:stop_index]
-
-            for chunk in [final_output[i:i+10] for i in range(0,len(final_output),10)]:
-                 for text in chunk.split():
-                    yield json.dumps({"text": text, "is_end": False}) + "\n"
+            
+            for text in final_output.split():
+                 yield json.dumps({"text": text, "is_end": False}) + "\n"
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
         else:
-            for chunk in [new_text[i:i+10] for i in range(0, len(new_text), 10)]:
-                 for text in chunk.split():
-                    yield json.dumps({"text": text, "is_end": False}) + "\n"
-
+            for text in new_text.split():
+              yield json.dumps({"text": text, "is_end": False}) + "\n"
         
-        if len(output_text) >= max_length:
-           
-            encoded_input = tokenizer(
-            output_text, return_tensors="pt",
-            truncation=True, max_length=max_length
-            ).to(device)
-           
-            output_text = ""
-
-        elif len(output_text) < max_length and len(new_text) == 0:
-            
-            for chunk in [output_text[i:i+10] for i in range(0, len(output_text), 10)]:
-                 for text in chunk.split():
-                    yield json.dumps({"text": text, "is_end": False}) + "\n"
-            
+        if len(new_text) == 0:
+            for text in output_text.split():
+                 yield json.dumps({"text": text, "is_end": False}) + "\n"
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
-            
+
+        encoded_input = tokenizer(
+            output_text, return_tensors="pt",
+            truncation=True
+        ).to(device)
+        output_text = ""
+
+
 
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
-
-        image_generator = pipeline(
-            "text-to-image", model=validated_body.model_name,
-            device=device
-        )
-        image = image_generator(validated_body.input_text)[0]
+        
+        if validated_body.model_name not in model_cache:
+            model = pipeline(
+                "text-to-image", model=validated_body.model_name,
+                device=device
+            )
+            model_cache[validated_body.model_name] = model
+        else:
+            model = model_cache[validated_body.model_name]
+        
+        image = model(validated_body.input_text)[0]
         
         image_data = list(image.getdata())
         
@@ -302,10 +291,15 @@ async def generate_text_to_speech(request: GenerateRequest):
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        audio_generator = pipeline(
-            "text-to-speech", model=validated_body.model_name,
-            device=device
-        )
+        if validated_body.model_name not in model_cache:
+            audio_generator = pipeline(
+                "text-to-speech", model=validated_body.model_name,
+                device=device
+            )
+            model_cache[validated_body.model_name] = audio_generator
+        else:
+            audio_generator = model_cache[validated_body.model_name]
+
         audio = audio_generator(validated_body.input_text)
 
         
@@ -327,10 +321,15 @@ async def generate_video(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        video_generator = pipeline(
+        if validated_body.model_name not in model_cache:
+            video_generator = pipeline(
             "text-to-video", model=validated_body.model_name,
             device=device
         )
+            model_cache[validated_body.model_name] = video_generator
+        else:
+            video_generator = model_cache[validated_body.model_name]
+
         video = video_generator(validated_body.input_text)