|
import os |
|
import logging |
|
import time |
|
from io import BytesIO |
|
from typing import Union |
|
|
|
from fastapi import FastAPI, HTTPException, Response, Request, UploadFile, File |
|
from fastapi.responses import StreamingResponse |
|
from pydantic import BaseModel, ValidationError, validator |
|
from transformers import ( |
|
AutoConfig, |
|
AutoModelForCausalLM, |
|
AutoTokenizer, |
|
pipeline, |
|
GenerationConfig, |
|
StoppingCriteriaList |
|
) |
|
import boto3 |
|
from huggingface_hub import hf_hub_download |
|
import soundfile as sf |
|
import numpy as np |
|
import torch |
|
import uvicorn |
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s") |
|
|
|
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") |
|
|
|
class GenerateRequest(BaseModel): |
|
model_name: str |
|
input_text: str = "" |
|
task_type: str |
|
temperature: float = 1.0 |
|
max_new_tokens: int = 200 |
|
stream: bool = False |
|
top_p: float = 1.0 |
|
top_k: int = 50 |
|
repetition_penalty: float = 1.0 |
|
num_return_sequences: int = 1 |
|
do_sample: bool = True |
|
chunk_delay: float = 0.0 |
|
stop_sequences: list[str] = [] |
|
|
|
@validator("model_name") |
|
def model_name_cannot_be_empty(cls, v): |
|
if not v: |
|
raise ValueError("model_name cannot be empty.") |
|
return v |
|
|
|
@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 |
|
|
|
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): |
|
s3_uri = self._get_s3_uri(model_name) |
|
try: |
|
logging.info(f"Trying to load {model_name} from S3...") |
|
config = AutoConfig.from_pretrained(s3_uri) |
|
model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config) |
|
tokenizer = AutoTokenizer.from_pretrained(s3_uri) |
|
|
|
if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None: |
|
if config.pad_token_id is not None: |
|
tokenizer.pad_token_id = config.pad_token_id |
|
else: |
|
tokenizer.pad_token_id = 0 |
|
|
|
logging.info(f"Loaded {model_name} from S3 successfully.") |
|
return model, tokenizer |
|
except EnvironmentError: |
|
logging.info(f"Model {model_name} not found in S3. Downloading...") |
|
try: |
|
model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN) |
|
tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN) |
|
|
|
if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None: |
|
config = AutoConfig.from_pretrained(model_name) |
|
if config.pad_token_id is not None: |
|
tokenizer.pad_token_id = config.pad_token_id |
|
else: |
|
tokenizer.pad_token_id = 0 |
|
|
|
logging.info(f"Downloaded {model_name} successfully.") |
|
logging.info(f"Saving {model_name} to S3...") |
|
model.save_pretrained(s3_uri) |
|
tokenizer.save_pretrained(s3_uri) |
|
logging.info(f"Saved {model_name} to S3 successfully.") |
|
return model, tokenizer |
|
except Exception as e: |
|
logging.exception(f"Error downloading/uploading model: {e}") |
|
raise HTTPException(status_code=500, detail=f"Error loading model: {e}") |
|
|
|
app = FastAPI() |
|
|
|
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) |
|
model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client) |
|
|
|
@app.post("/generate") |
|
async def generate(request: Request, body: GenerateRequest): |
|
try: |
|
validated_body = GenerateRequest(**body.model_dump()) |
|
model, tokenizer = await model_loader.load_model_and_tokenizer(validated_body.model_name) |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
model.to(device) |
|
|
|
if validated_body.task_type == "text-to-text": |
|
generation_config = GenerationConfig( |
|
temperature=validated_body.temperature, |
|
max_new_tokens=validated_body.max_new_tokens, |
|
top_p=validated_body.top_p, |
|
top_k=validated_body.top_k, |
|
repetition_penalty=validated_body.repetition_penalty, |
|
do_sample=validated_body.do_sample, |
|
num_return_sequences=validated_body.num_return_sequences, |
|
stopping_criteria=StoppingCriteriaList( |
|
[lambda _, outputs: tokenizer.decode(outputs[0][-1]) in validated_body.stop_sequences] if validated_body.stop_sequences else [] |
|
) |
|
) |
|
|
|
async def stream_text(): |
|
input_text = validated_body.input_text |
|
generated_text = "" |
|
max_length = model.config.max_position_embeddings |
|
|
|
while True: |
|
encoded_input = tokenizer(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 |
|
|
|
if remaining_tokens <= 0: |
|
break |
|
|
|
generation_config.max_new_tokens = min(remaining_tokens, validated_body.max_new_tokens) |
|
|
|
output = model.generate(**encoded_input, generation_config=generation_config) |
|
chunk = tokenizer.decode(output[0], skip_special_tokens=True) |
|
generated_text += chunk |
|
yield chunk |
|
time.sleep(validated_body.chunk_delay) |
|
input_text = generated_text |
|
|
|
if validated_body.stream: |
|
return StreamingResponse(stream_text(), media_type="text/plain") |
|
else: |
|
generated_text = "" |
|
async for chunk in stream_text(): |
|
generated_text += chunk |
|
return {"result": generated_text} |
|
|
|
elif validated_body.task_type == "text-to-image": |
|
generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=device) |
|
image = generator(validated_body.input_text)[0] |
|
image_bytes = image.tobytes() |
|
return Response(content=image_bytes, media_type="image/png") |
|
|
|
elif validated_body.task_type == "text-to-speech": |
|
generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device) |
|
audio = generator(validated_body.input_text) |
|
audio_bytesio = BytesIO() |
|
sf.write(audio_bytesio, audio["sampling_rate"], np.int16(audio["audio"])) |
|
audio_bytes = audio_bytesio.getvalue() |
|
return Response(content=audio_bytes, media_type="audio/wav") |
|
|
|
elif validated_body.task_type == "text-to-video": |
|
try: |
|
generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device) |
|
video = generator(validated_body.input_text) |
|
return Response(content=video, media_type="video/mp4") |
|
except Exception as e: |
|
raise HTTPException(status_code=500, detail=f"Error in text-to-video generation: {e}") |
|
|
|
else: |
|
raise HTTPException(status_code=400, detail="Unsupported task type") |
|
|
|
except HTTPException as e: |
|
raise e |
|
except ValidationError as e: |
|
raise HTTPException(status_code=422, detail=e.errors()) |
|
except Exception as e: |
|
logging.exception(f"An unexpected error occurred: {e}") |
|
raise HTTPException(status_code=500, detail="An unexpected error occurred.") |
|
|
|
|
|
if __name__ == "__main__": |
|
uvicorn.run(app, host="0.0.0.0", port=7860) |
