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:app", host="0.0.0.0", port=8000, reload=True)