Update app.py

app.py

@@ -1,13 +1,13 @@
 import os
-import torch
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi import FastAPI, HTTPException, Depends
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel, field_validator
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList, pipeline
 import boto3
 import uvicorn
-from io import BytesIO
-from transformers import pipeline
+import soundfile as sf
+import imageio
+from typing import Dict
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -15,6 +15,9 @@ AWS_REGION = os.getenv("AWS_REGION")
 S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
 HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
+if not all([AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION, S3_BUCKET_NAME]):
+    raise ValueError("Missing one or more AWS environment variables.")
+
 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()
@@ -39,6 +42,7 @@ class GenerateRequest(BaseModel):
     do_sample: bool = True
     stop_sequences: list[str] = []
     no_repeat_ngram_size: int = 2
+    continuation_id: str = None
 
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
@@ -70,37 +74,33 @@ class S3ModelLoader:
     async def load_model_and_tokenizer(self, model_name):
         s3_uri = self._get_s3_uri(model_name)
         try:
-            config = AutoConfig.from_pretrained(s3_uri, local_files_only=True)
-            model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=True)
-            tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config, local_files_only=True)
+            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)
             tokenizer.add_special_tokens(SPECIAL_TOKENS)
             model.resize_token_embeddings(len(tokenizer))
             if tokenizer.pad_token_id is None:
                 tokenizer.pad_token_id = tokenizer.eos_token_id
             return model, tokenizer
-        except EnvironmentError:
-            try:
-                config = AutoConfig.from_pretrained(model_name)
-                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
-                tokenizer.add_special_tokens(SPECIAL_TOKENS)
-                model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
-                model.resize_token_embeddings(len(tokenizer))
-                if tokenizer.pad_token_id is None:
-                    tokenizer.pad_token_id = tokenizer.eos_token_id
-                model.save_pretrained(s3_uri)
-                tokenizer.save_pretrained(s3_uri)
-                return model, tokenizer
-            except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error loading model from S3: {e}")
 
 model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 
+active_generations: Dict[str, Dict] = {}
+
+async def get_model_and_tokenizer(model_name: str):
+    try:
+        return await model_loader.load_model_and_tokenizer(model_name)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
+
 @app.post("/generate")
-async def generate(request: GenerateRequest):
+async def generate(request: GenerateRequest, model_resources: tuple = Depends(get_model_and_tokenizer)):
+    model, tokenizer = model_resources
     try:
         model_name = request.model_name
         input_text = request.input_text
-        task_type = request.task_type
         temperature = request.temperature
         max_new_tokens = request.max_new_tokens
         top_p = request.top_p
@@ -110,10 +110,13 @@ async def generate(request: GenerateRequest):
         do_sample = request.do_sample
         stop_sequences = request.stop_sequences
         no_repeat_ngram_size = request.no_repeat_ngram_size
+        continuation_id = request.continuation_id
 
-        model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        model.to(device)
+        if continuation_id:
+            if continuation_id not in active_generations:
+                raise HTTPException(status_code=404, detail="Continuation ID not found.")
+            previous_output = active_generations[continuation_id]["output"]
+            input_text = previous_output
 
         generation_config = GenerationConfig(
             temperature=temperature,
@@ -127,15 +130,24 @@ async def generate(request: GenerateRequest):
             pad_token_id=tokenizer.pad_token_id
         )
 
-        generated_text = generate_text(model, tokenizer, input_text, generation_config, stop_sequences, device)
-        return JSONResponse({"text": generated_text})
+        generated_text = generate_text_internal(model, tokenizer, input_text, generation_config, stop_sequences)
+
+        if not continuation_id:
+            continuation_id = os.urandom(16).hex()
+            active_generations[continuation_id] = {"model_name": model_name, "output": generated_text}
+        else:
+            active_generations[continuation_id]["output"] = generated_text
 
+        return JSONResponse({"text": generated_text, "continuation_id": continuation_id})
+
+    except HTTPException as http_err:
+        raise http_err
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-def generate_text(model, tokenizer, input_text, generation_config, stop_sequences, device):
+def generate_text_internal(model, tokenizer, input_text, generation_config, stop_sequences):
     max_model_length = model.config.max_position_embeddings
-    encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True).to(device)
+    encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True)
 
     stopping_criteria = StoppingCriteriaList()
 
@@ -159,62 +171,61 @@ def generate_text(model, tokenizer, input_text, generation_config, stop_sequence
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return generated_text
 
+async def load_pipeline_from_s3(task, model_name):
+    s3_uri = f"s3://{S3_BUCKET_NAME}/{model_name.replace('/', '-')}"
+    try:
+        return pipeline(task, model=s3_uri)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error loading {task} model from S3: {e}")
+
 @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]
-
-        img_byte_arr = BytesIO()
-        image.save(img_byte_arr, format="PNG")
-        img_byte_arr.seek(0)
+        if request.task_type != "text-to-image":
+            raise HTTPException(status_code=400, detail="Invalid task_type for this endpoint.")
 
-        return StreamingResponse(img_byte_arr, media_type="image/png")
+        image_generator = await load_pipeline_from_s3("text-to-image", request.model_name)
+        image = image_generator(request.input_text)[0]
+        continuation_id = os.urandom(16).hex()
+        active_generations[continuation_id] = {"model_name": request.model_name, "output": "Image generated successfully"}
+        return JSONResponse({"url": "Image generated successfully", "continuation_id": continuation_id})
 
+    except HTTPException as http_err:
+        raise http_err
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 @app.post("/generate-text-to-speech")
 async def generate_text_to_speech(request: GenerateRequest):
     try:
-        validated_body = request
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        if request.task_type != "text-to-speech":
+            raise HTTPException(status_code=400, detail="Invalid task_type for this endpoint.")
 
-        audio_generator = pipeline("text-to-speech", model=validated_body.model_name, device=device)
-        audio = audio_generator(validated_body.input_text)[0]
-
-        audio_byte_arr = BytesIO()
-        # It is expected that the audio is saved as wav.
-        # Saving like this will not always work. Please check how your
-        # audio_generator model is working.
-        audio_generator.save_audio(audio_byte_arr, audio)
-        audio_byte_arr.seek(0)
-
-        return StreamingResponse(audio_byte_arr, media_type="audio/wav")
+        tts_pipeline = await load_pipeline_from_s3("text-to-speech", request.model_name)
+        output = tts_pipeline(request.input_text)
+        continuation_id = os.urandom(16).hex()
+        active_generations[continuation_id] = {"model_name": request.model_name, "output": "Audio generated successfully"}
+        return JSONResponse({"url": "Audio generated successfully", "continuation_id": continuation_id})
 
+    except HTTPException as http_err:
+        raise http_err
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 @app.post("/generate-video")
 async def generate_video(request: GenerateRequest):
     try:
-        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)
-        video = video_generator(validated_body.input_text)[0]
-
-        video_byte_arr = BytesIO()
-        # Same as above. Please check how your video model is returning the
-        # videos and save them accordingly.
-        # It is expected that the video is saved as MP4
-        video_generator.save_video(video_byte_arr, video)
-        video_byte_arr.seek(0)
+        if request.task_type != "text-to-video":
+            raise HTTPException(status_code=400, detail="Invalid task_type for this endpoint.")
 
-        return StreamingResponse(video_byte_arr, media_type="video/mp4")
+        video_pipeline = await load_pipeline_from_s3("text-to-video", request.model_name)
+        output = video_pipeline(request.input_text)
+        continuation_id = os.urandom(16).hex()
+        active_generations[continuation_id] = {"model_name": request.model_name, "output": "Video generated successfully"}
+        return JSONResponse({"url": "Video generated successfully", "continuation_id": continuation_id})
 
+    except HTTPException as http_err:
+        raise http_err
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")