Update app.py

app.py

@@ -17,8 +17,8 @@ from transformers import pipeline
 import json
 from huggingface_hub import login
 import base64
-import io
-from PIL import Image
+from botocore.exceptions import NoCredentialsError
+
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -26,23 +26,17 @@ 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 = ""
@@ -71,39 +65,37 @@ 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
         self.s3_client = s3_client
 
     def _get_s3_uri(self, model_name):
-        return f"s3://{self.bucket_name}/" \
+         return f"s3://{self.bucket_name}/" \
                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
             )
 
-            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
+            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, tokenizer
-        except EnvironmentError:
+        except (EnvironmentError, NoCredentialsError):
             try:
                 config = AutoConfig.from_pretrained(
                     model_name, token=HUGGINGFACE_HUB_TOKEN
@@ -111,29 +103,42 @@ 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 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
 
+                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.save_pretrained(s3_uri)
                 tokenizer.save_pretrained(s3_uri)
+                
+                
+                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
+                )
+                
                 return model, tokenizer
             except Exception as e:
                 raise HTTPException(
                     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:
@@ -141,7 +146,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
+        max_new_tokens = request.max_new_tokens  # This value will be used to constraint the output
         stream = request.stream
         top_p = request.top_p
         top_k = request.top_k
@@ -153,7 +158,7 @@ async def generate(request: GenerateRequest):
         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,
@@ -166,21 +171,13 @@ 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:
-            raise HTTPException(
-                status_code=400, detail="Task type not text-to-text"
-            )
+            return HTTPException(status_code=400, detail="Task type not text-to-text")
 
     except Exception as e:
         raise HTTPException(
@@ -188,11 +185,12 @@ 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
@@ -203,12 +201,14 @@ async def stream_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 = ""
@@ -229,7 +229,7 @@ async def stream_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
@@ -238,39 +238,33 @@ async def stream_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)
-            ]
-
-            for chunk in chunked_output:
-                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
 
+            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"
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
 
         else:
-            chunked_output = [
-                new_text[i: i + 10] for i in range(0, len(new_text), 10)
-            ]
+            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 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)
-            ]
 
-            for chunk in chunked_output:
-                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
+            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"
+
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             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:
@@ -278,27 +272,19 @@ 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]
-
-        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")
+        
+        image_data = list(image.getdata())
+        
+        return json.dumps({"image_data": image_data, "is_end": True})
 
     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)}"
         )
 
 
@@ -309,25 +295,22 @@ 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"]
 
-        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")
+        
+        audio_bytes = audio["audio"]
+        
+        audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
+        
+        return json.dumps({"audio": audio_base64, "is_end": True})
 
     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)}"
         )
 
 
@@ -337,26 +320,21 @@ 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)
-
-        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")
-
+        
+        
+        video_base64 = base64.b64encode(video).decode('utf-8')
+        
+        return json.dumps({"video": video_base64, "is_end": True})
 
     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)