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)