Update app.py

app.py

@@ -17,7 +17,8 @@ from transformers import pipeline
 import json
 from huggingface_hub import login
 import base64
-
+import io
+from PIL import Image
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -25,23 +26,29 @@ 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)
+    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)
+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()
 
+# Global variables for tokenizer tokens
+EOS_TOKEN_ID = None
+PAD_TOKEN_ID = None
+
+
 class GenerateRequest(BaseModel):
     model_name: str
     input_text: str = ""
     task_type: str
     temperature: float = 1.0
-    max_new_tokens: int = 200 # this will be limited to 10
+    max_new_tokens: int = 200  # this will be limited to 10
     stream: bool = True
     top_p: float = 1.0
     top_k: int = 50
@@ -64,6 +71,7 @@ class GenerateRequest(BaseModel):
             raise ValueError(f"task_type must be one of: {valid_types}")
         return v
 
+
 class S3ModelLoader:
     def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
@@ -74,24 +82,25 @@ class S3ModelLoader:
                f"{model_name.replace('/', '-')}"
 
     async def load_model_and_tokenizer(self, model_name):
+        global EOS_TOKEN_ID, PAD_TOKEN_ID
         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
+            EOS_TOKEN_ID = tokenizer.eos_token_id
+            PAD_TOKEN_ID = tokenizer.pad_token_id
+
+            if EOS_TOKEN_ID is not None and PAD_TOKEN_ID is None:
+                 PAD_TOKEN_ID = config.pad_token_id or EOS_TOKEN_ID
+                 tokenizer.pad_token_id = PAD_TOKEN_ID
 
             return model, tokenizer
         except EnvironmentError:
@@ -102,16 +111,16 @@ class S3ModelLoader:
                 tokenizer = AutoTokenizer.from_pretrained(
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
-                
+
                 model = AutoModelForCausalLM.from_pretrained(
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
+                EOS_TOKEN_ID = tokenizer.eos_token_id
+                PAD_TOKEN_ID = tokenizer.pad_token_id
 
-
-                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
+                if EOS_TOKEN_ID is not None and PAD_TOKEN_ID is None:
+                  PAD_TOKEN_ID = config.pad_token_id or EOS_TOKEN_ID
+                  tokenizer.pad_token_id = PAD_TOKEN_ID
 
                 model.save_pretrained(s3_uri)
                 tokenizer.save_pretrained(s3_uri)
@@ -121,8 +130,10 @@ class S3ModelLoader:
                     status_code=500, detail=f"Error loading model: {e}"
                 )
 
+
 model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 
+
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     try:
@@ -130,7 +141,7 @@ async def generate(request: GenerateRequest):
         input_text = request.input_text
         task_type = request.task_type
         temperature = request.temperature
-        max_new_tokens = request.max_new_tokens #This value will be used to constraint the output
+        max_new_tokens = request.max_new_tokens
         stream = request.stream
         top_p = request.top_p
         top_k = request.top_k
@@ -139,15 +150,14 @@ 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)
+        model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
-        
+
         if "text-to-text" == task_type:
             generation_config = GenerationConfig(
                 temperature=temperature,
-                max_new_tokens=min(max_new_tokens,10),  # Constrain max_new_tokens to 10
+                max_new_tokens=min(max_new_tokens, 10),  # Constrain max_new_tokens to 10
                 top_p=top_p,
                 top_k=top_k,
                 repetition_penalty=repetition_penalty,
@@ -156,13 +166,21 @@ async def generate(request: GenerateRequest):
             )
 
             return StreamingResponse(
-                stream_text(model, tokenizer, input_text,
-                             generation_config, stop_sequences,
-                             device, max_length=10),
-                media_type="text/plain"
+                stream_text(
+                    model,
+                    tokenizer,
+                    input_text,
+                    generation_config,
+                    stop_sequences,
+                    device,
+                    max_length=10,
+                ),
+                media_type="text/plain",
             )
         else:
-            return HTTPException(status_code=400, detail="Task type not text-to-text")
+            raise HTTPException(
+                status_code=400, detail="Task type not text-to-text"
+            )
 
     except Exception as e:
         raise HTTPException(
@@ -170,12 +188,11 @@ async def generate(request: GenerateRequest):
         )
 
 
-async def stream_text(model, tokenizer, input_text,
-                        generation_config, stop_sequences,
-                        device, max_length):
+async def stream_text(
+    model, tokenizer, input_text, generation_config, stop_sequences, device, max_length
+):
     encoded_input = tokenizer(
-        input_text, return_tensors="pt",
-        truncation=True, max_length=max_length
+        input_text, return_tensors="pt", truncation=True, max_length=max_length
     ).to(device)
     input_length = encoded_input["input_ids"].shape[1]
     remaining_tokens = max_length - input_length
@@ -186,14 +203,12 @@ async def stream_text(model, tokenizer, input_text,
     generation_config.max_new_tokens = min(
         remaining_tokens, generation_config.max_new_tokens
     )
-    
 
     def find_stop(output_text, stop_sequences):
         for seq in stop_sequences:
             if seq in output_text:
                 last_index = output_text.rfind(seq)
                 return last_index + len(seq)
-
         return -1
 
     output_text = ""
@@ -214,7 +229,7 @@ async def stream_text(model, tokenizer, input_text,
 
         new_text = tokenizer.decode(
             outputs.sequences[0][len(encoded_input["input_ids"][0]):],
-            skip_special_tokens=True
+            skip_special_tokens=True,
         )
 
         output_text += new_text
@@ -223,8 +238,9 @@ async def stream_text(model, tokenizer, input_text,
 
         if stop_index != -1:
             final_output = output_text[:stop_index]
-            chunked_output = [final_output[i:i+10]
-                             for i in range(0, len(final_output), 10)]
+            chunked_output = [
+                final_output[i: i + 10] for i in range(0, len(final_output), 10)
+            ]
 
             for chunk in chunked_output:
                 yield json.dumps({"text": chunk, "is_end": False}) + "\n"
@@ -233,15 +249,17 @@ async def stream_text(model, tokenizer, input_text,
             break
 
         else:
-             chunked_output = [new_text[i:i+10]
-                             for i in range(0, len(new_text), 10)]
+            chunked_output = [
+                new_text[i: i + 10] for i in range(0, len(new_text), 10)
+            ]
 
-             for chunk in chunked_output:
-                  yield json.dumps({"text": chunk, "is_end": False}) + "\n"
+            for chunk in chunked_output:
+                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
 
         if len(output_text) >= generation_config.max_new_tokens:
-            chunked_output = [output_text[i:i+10]
-                             for i in range(0, len(output_text), 10)]
+            chunked_output = [
+                output_text[i: i + 10] for i in range(0, len(output_text), 10)
+            ]
 
             for chunk in chunked_output:
                 yield json.dumps({"text": chunk, "is_end": False}) + "\n"
@@ -249,10 +267,10 @@ async def stream_text(model, tokenizer, input_text,
             break
 
         encoded_input = tokenizer(
-            output_text, return_tensors="pt",
-            truncation=True, max_length=max_length
+            output_text, return_tensors="pt", truncation=True, max_length=max_length
         ).to(device)
 
+
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try:
@@ -260,19 +278,27 @@ async def generate_image(request: GenerateRequest):
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
         image_generator = pipeline(
-            "text-to-image", model=validated_body.model_name,
-            device=device
+            "text-to-image", model=validated_body.model_name, device=device
         )
         image = image_generator(validated_body.input_text)[0]
-        
-        image_data = list(image.getdata())
-        
-        return json.dumps({"image_data": image_data, "is_end": True})
+
+        async def stream_image():
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+            image_bytes = buffered.getvalue()
+            image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+            chunk_size = 1000
+            for i in range(0, len(image_base64), chunk_size):
+                chunk = image_base64[i: i + chunk_size]
+                yield json.dumps({"image": chunk, "is_end": False}) + "\n"
+
+            yield json.dumps({"image": "", "is_end": True}) + "\n"
+
+        return StreamingResponse(stream_image(), media_type="text/plain")
 
     except Exception as e:
         raise HTTPException(
-            status_code=500,
-            detail=f"Internal server error: {str(e)}"
+            status_code=500, detail=f"Internal server error: {str(e)}"
         )
 
 
@@ -283,22 +309,25 @@ async def generate_text_to_speech(request: GenerateRequest):
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
         audio_generator = pipeline(
-            "text-to-speech", model=validated_body.model_name,
-            device=device
+            "text-to-speech", model=validated_body.model_name, device=device
         )
         audio = audio_generator(validated_body.input_text)
-
-        
         audio_bytes = audio["audio"]
-        
-        audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
-        
-        return json.dumps({"audio": audio_base64, "is_end": True})
+
+        async def stream_audio():
+            audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
+            chunk_size = 1000
+            for i in range(0, len(audio_base64), chunk_size):
+                chunk = audio_base64[i: i + chunk_size]
+                yield json.dumps({"audio": chunk, "is_end": False}) + "\n"
+
+            yield json.dumps({"audio": "", "is_end": True}) + "\n"
+
+        return StreamingResponse(stream_audio(), media_type="text/plain")
 
     except Exception as e:
         raise HTTPException(
-            status_code=500,
-            detail=f"Internal server error: {str(e)}"
+            status_code=500, detail=f"Internal server error: {str(e)}"
         )
 
 
@@ -308,21 +337,26 @@ async def generate_video(request: GenerateRequest):
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
         video_generator = pipeline(
-            "text-to-video", model=validated_body.model_name,
-            device=device
+            "text-to-video", model=validated_body.model_name, device=device
         )
         video = video_generator(validated_body.input_text)
-        
-        
-        video_base64 = base64.b64encode(video).decode('utf-8')
-        
-        return json.dumps({"video": video_base64, "is_end": True})
+
+        async def stream_video():
+            video_base64 = base64.b64encode(video).decode("utf-8")
+            chunk_size = 1000
+            for i in range(0, len(video_base64), chunk_size):
+                chunk = video_base64[i: i + chunk_size]
+                yield json.dumps({"video": chunk, "is_end": False}) + "\n"
+
+            yield json.dumps({"video": "", "is_end": True}) + "\n"
+        return StreamingResponse(stream_video(), media_type="text/plain")
+
 
     except Exception as e:
         raise HTTPException(
-            status_code=500,
-            detail=f"Internal server error: {str(e)}"
+            status_code=500, detail=f"Internal server error: {str(e)}"
         )
 
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)