|
import os |
|
import logging |
|
import threading |
|
import boto3 |
|
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, StoppingCriteriaList, AutoConfig |
|
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 |
|
|
|
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}") |
|
|
|
async def load_model_and_tokenizer(self, model_name): |
|
try: |
|
s3_model_path = await asyncio.to_thread(self._download_from_s3, model_name) |
|
|
|
config = AutoConfig.from_pretrained(s3_model_path) |
|
tokenizer = AutoTokenizer.from_pretrained(s3_model_path, config=config) |
|
model = AutoModelForCausalLM.from_pretrained(s3_model_path, config=config) |
|
|
|
logging.info(f"Model {model_name} loaded successfully from S3.") |
|
return model, tokenizer |
|
except Exception as e: |
|
logging.exception(f"Error loading model: {e}") |
|
raise HTTPException(status_code=500, detail=f"Error loading model: {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, filename=model_name.split("/")[-1]) |
|
self.s3_client.upload_file(model_dir, self.bucket_name, f"lilmeaty_garca/{model_name}") |
|
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}") |
|
|
|
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() |
|
|
|
@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["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 generating video: {str(e)}") |
|
|
|
else: |
|
raise HTTPException(status_code=400, detail="Invalid task type.") |
|
|
|
except Exception as e: |
|
logging.error(f"Error processing request: {str(e)}") |
|
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") |
|
|
|
if __name__ == "__main__": |
|
import uvicorn |
|
uvicorn.run(app, host="0.0.0.0", port=8000) |
|
|
