import os
import logging
import threading
import boto3
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, StoppingCriteriaList, pipeline
from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel, field_validator
from huggingface_hub import hf_hub_download
import requests
import time
import asyncio
from fastapi.responses import StreamingResponse, Response
import torch
from io import BytesIO
import numpy as np
import soundfile as sf
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
app = FastAPI()
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] = []
@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
def _get_s3_uri(self, model_name):
return f"s3://{self.bucket_name}/lilmeaty_garca/{model_name.replace('/', '-')}"
def _download_from_s3(self, model_name):
try:
logging.info(f"Attempting to load model {model_name} from S3...")
model_files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=f"lilmeaty_garca/{model_name}")
if "Contents" not in model_files:
raise FileNotFoundError(f"Model files not found in S3 for {model_name}")
s3_model_path = f"s3://{self.bucket_name}/lilmeaty_garca/{model_name.replace('/', '-')}"
logging.info(f"Model {model_name} found on S3 at {s3_model_path}")
return s3_model_path
except Exception as e:
logging.error(f"Error downloading from S3: {e}")
raise HTTPException(status_code=500, detail=f"Error downloading model from S3: {e}")
def download_model_from_huggingface(self, model_name):
try:
logging.info(f"Downloading model {model_name} from Hugging Face...")
model_dir = hf_hub_download(model_name, token=HUGGINGFACE_HUB_TOKEN)
model_files = os.listdir(model_dir)
for model_file in model_files:
s3_path = f"lilmeaty_garca/{model_name}/{model_file}"
self.s3_client.upload_file(os.path.join(model_dir, model_file), self.bucket_name, s3_path)
logging.info(f"Model {model_name} saved to S3 successfully.")
except Exception as e:
logging.error(f"Error downloading model {model_name} from Hugging Face: {e}")
raise HTTPException(status_code=500, detail=f"Error downloading model from Hugging Face: {e}")
def download_all_models_in_background(self):
models_url = "https://huggingface.co/api/models"
try:
response = requests.get(models_url)
if response.status_code != 200:
logging.error("Error getting Hugging Face model list.")
raise HTTPException(status_code=500, detail="Error getting model list.")
models = response.json()
for model in models:
model_name = model["id"]
self.download_model_from_huggingface(model_name)
except Exception as e:
logging.error(f"Error downloading models in the background: {e}")
raise HTTPException(status_code=500, detail="Error downloading models in the background.")
def run_in_background(self):
threading.Thread(target=self.download_all_models_in_background, daemon=True).start()
def load_model_and_tokenizer(self, model_name):
try:
model_uri = self._download_from_s3(model_name)
model = AutoModelForCausalLM.from_pretrained(model_uri)
tokenizer = AutoTokenizer.from_pretrained(model_uri)
logging.info(f"Model {model_name} loaded successfully from {model_uri}.")
return model, tokenizer
except Exception as e:
logging.error(f"Error loading model {model_name}: {e}")
raise HTTPException(status_code=500, detail=f"Error loading model {model_name}: {e}")
@app.on_event("startup")
async def startup_event():
model_loader.run_in_background()
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
)
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)
stopping_criteria = StoppingCriteriaList(
[lambda _, outputs: tokenizer.decode(outputs[0][-1], skip_special_tokens=True) in validated_body.stop_sequences] if validated_body.stop_sequences else []
)
output = model.generate(**encoded_input, generation_config=generation_config, stopping_criteria=stopping_criteria)
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["samples"], audio["rate"], format="WAV")
audio_bytesio.seek(0)
return StreamingResponse(audio_bytesio, media_type="audio/wav")
elif validated_body.task_type == "text-to-video":
return {"error": "Text-to-video task type is not yet supported."}
else:
raise HTTPException(status_code=400, detail="Invalid task type")
except Exception as e:
logging.error(f"Error during generation: {e}")
raise HTTPException(status_code=500, detail=f"Internal Server Error: {e}")
import uvicorn
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000, reload=True)