app.py

import os
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, field_validator
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteria, StoppingCriteriaList, pipeline
import boto3
import uvicorn
import asyncio
import json
from huggingface_hub import login
from botocore.exceptions import NoCredentialsError
from functools import cached_property
import base64
from optimum.onnxruntime import ORTModelForCausalLM
from optimum.bettertransformer import BetterTransformer
import bitsandbytes as bnb

AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
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)

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()

class GenerateRequest(BaseModel):
    model_name: str
    input_text: str = ""
    task_type: str
    temperature: float = 0.01
    max_new_tokens: int = 20
    stream: bool = True
    top_p: float = 1.0
    top_k: int = 1
    repetition_penalty: float = 1.0
    num_return_sequences: int = 1
    do_sample: bool = False
    stop_sequences: list[str] = []
    quantize: bool = True
    use_onnx: bool = False
    use_bettertransformer: bool = True
    @field_validator("model_name")
    def model_name_cannot_be_empty(cls, v):
        if not v:
            raise ValueError("model_name cannot be empty.")
        return v
    @field_validator("task_type")
    def task_type_must_be_valid(cls, v):
        valid_types = ["text-to-text", "text-to-image","text-to-speech", "text-to-video"]
        if v not in valid_types:
            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
        self.model_cache = {}
    def _get_s3_uri(self, model_name):
         return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
    async def _load_model_and_tokenizer(self, model_name, quantize, use_onnx, use_bettertransformer):
       s3_uri = self._get_s3_uri(model_name)
       try:
            config = AutoConfig.from_pretrained(s3_uri, local_files_only=False)
            if use_onnx:
                model = ORTModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=False).to(self.device)
            elif quantize:
                model = AutoModelForCausalLM.from_pretrained(
                    s3_uri, config=config, local_files_only=False,
                    load_in_8bit=True
                    ).to(self.device)
            else:
                model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=False).to(self.device)
            if use_bettertransformer:
                model = BetterTransformer.transform(model)
            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
            return model, tokenizer
       except (EnvironmentError, NoCredentialsError):
            try:
                config = AutoConfig.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN)
                if use_onnx:
                    model = ORTModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN).to(self.device)
                elif quantize:
                    model = AutoModelForCausalLM.from_pretrained(
                        model_name, config=config, token=HUGGINGFACE_HUB_TOKEN,
                        load_in_8bit=True
                    ).to(self.device)
                else:
                     model = AutoModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN).to(self.device)
                if use_bettertransformer:
                    model = BetterTransformer.transform(model)
                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 Exception as e:
                raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
    @cached_property
    def device(self):
        return torch.device("cpu")
    async def get_model_and_tokenizer(self, model_name, quantize, use_onnx, use_bettertransformer):
        key = f"{model_name}-{quantize}-{use_onnx}-{use_bettertransformer}"
        if key not in self.model_cache:
            model, tokenizer = await self._load_model_and_tokenizer(model_name, quantize, use_onnx, use_bettertransformer)
            self.model_cache[key] = {"model":model, "tokenizer":tokenizer}
        return self.model_cache[key]["model"], self.model_cache[key]["tokenizer"]
    async def get_pipeline(self, model_name, task_type):
        if model_name not in self.model_cache:
             config = AutoConfig.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
             model = pipeline(task_type, model=model_name,device=self.device, config=config)
             self.model_cache[model_name] = {"model":model}
        return self.model_cache[model_name]["model"]
model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)

@app.post("/generate")
async def generate(request: GenerateRequest):
    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
        stream = request.stream
        top_p = request.top_p
        top_k = request.top_k
        repetition_penalty = request.repetition_penalty
        num_return_sequences = request.num_return_sequences
        do_sample = request.do_sample
        stop_sequences = request.stop_sequences
        quantize = request.quantize
        use_onnx = request.use_onnx
        use_bettertransformer = request.use_bettertransformer
        model, tokenizer = await model_loader.get_model_and_tokenizer(model_name, quantize, use_onnx, use_bettertransformer)
        if "text-to-text" == task_type:
            generation_config = GenerationConfig(temperature=temperature,max_new_tokens=max_new_tokens,top_p=top_p,top_k=top_k,repetition_penalty=repetition_penalty,do_sample=do_sample,num_return_sequences=num_return_sequences,eos_token_id = tokenizer.eos_token_id)
            if stream:
                return StreamingResponse(stream_text(model, tokenizer, input_text,generation_config, stop_sequences),media_type="text/plain")
            else:
                result = await generate_text(model, tokenizer, input_text,generation_config, stop_sequences)
                return JSONResponse({"text": result, "is_end": True})
        else:
            return HTTPException(status_code=400, detail="Task type not text-to-text")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
class StopOnSequences(StoppingCriteria):
    def __init__(self, stop_sequences, tokenizer):
        self.stop_sequences = stop_sequences
        self.tokenizer = tokenizer
        self.stop_ids = [tokenizer.encode(seq, add_special_tokens=False) for seq in stop_sequences]
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        decoded_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
        for stop_sequence in self.stop_sequences:
             if stop_sequence in decoded_text:
                 return True
        return False
async def stream_text(model, tokenizer, input_text,generation_config, stop_sequences):
    encoded_input = tokenizer(input_text, return_tensors="pt",truncation=True).to(model_loader.device)
    stop_criteria = StopOnSequences(stop_sequences, tokenizer)
    stopping_criteria = StoppingCriteriaList([stop_criteria])
    async for token in _stream_text(model, encoded_input, tokenizer, generation_config, stop_criteria, stopping_criteria):
            yield json.dumps({"text":token, "is_end": False}) + "\n"
    yield json.dumps({"text":"", "is_end": True}) + "\n"
async def _stream_text(model, encoded_input, tokenizer, generation_config, stop_criteria, stopping_criteria):
    output_text = ""
    while True:
         outputs = await asyncio.to_thread(model.generate,**encoded_input,do_sample=generation_config.do_sample,max_new_tokens=generation_config.max_new_tokens,temperature=generation_config.temperature,top_p=generation_config.top_p,top_k=generation_config.top_k,repetition_penalty=generation_config.repetition_penalty,num_return_sequences=generation_config.num_return_sequences,output_scores=True,return_dict_in_generate=True,stopping_criteria=stopping_criteria)
         new_text = tokenizer.decode(outputs.sequences[0][len(encoded_input["input_ids"][0]):],skip_special_tokens=True)
         if len(new_text) == 0:
             if not stop_criteria(outputs.sequences, None):
                 for token in output_text.split():
                    yield token
             break
         output_text += new_text
         for token in new_text.split():
              yield token
         if stop_criteria(outputs.sequences, None):
             break
         encoded_input = tokenizer(output_text, return_tensors="pt",truncation=True).to(model_loader.device)
         output_text=""
async def generate_text(model, tokenizer, input_text,generation_config, stop_sequences):
    encoded_input = tokenizer(input_text, return_tensors="pt",truncation=True).to(model_loader.device)
    stop_criteria = StopOnSequences(stop_sequences, tokenizer)
    stopping_criteria = StoppingCriteriaList([stop_criteria])
    outputs = await asyncio.to_thread(model.generate,**encoded_input,do_sample=generation_config.do_sample,max_new_tokens=generation_config.max_new_tokens,temperature=generation_config.temperature,top_p=generation_config.top_p,top_k=generation_config.top_k,repetition_penalty=generation_config.repetition_penalty,num_return_sequences=num_return_sequences,output_scores=True,return_dict_in_generate=True,stopping_criteria=stopping_criteria)
    generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
    return generated_text
@app.post("/generate-image")
async def generate_image(request: GenerateRequest):
    try:
        validated_body = request
        model = await model_loader.get_pipeline(validated_body.model_name, "text-to-image")
        image = model(validated_body.input_text)[0]
        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)}")
@app.post("/generate-text-to-speech")
async def generate_text_to_speech(request: GenerateRequest):
    try:
        validated_body = request
        audio_generator = await model_loader.get_pipeline(validated_body.model_name, "text-to-speech")
        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})
    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
        video_generator = await model_loader.get_pipeline(validated_body.model_name, "text-to-video")
        video = video_generator(validated_body.input_text)
        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)}")
async def load_all_models():
    pass
if __name__ == "__main__":
    import asyncio
    asyncio.run(load_all_models())
    uvicorn.run(app, host="0.0.0.0", port=7860)