import os import logging import time import threading import asyncio from io import BytesIO import requests import boto3 import torch import safetensors import soundfile as sf import numpy as np from fastapi import FastAPI, HTTPException, Request, UploadFile, File from fastapi.responses import StreamingResponse from pydantic import BaseModel, ValidationError, field_validator from transformers import ( AutoConfig, AutoModelForCausalLM, AutoTokenizer, pipeline, GenerationConfig, StoppingCriteriaList ) from huggingface_hub import hf_hub_download import uvicorn 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): s3_uri = self._get_s3_uri(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}") local_dir = f"/tmp/{model_name.replace('/', '-')}" os.makedirs(local_dir, exist_ok=True) for obj in model_files["Contents"]: file_key = obj["Key"] if file_key.endswith('/'): continue local_file_path = os.path.join(local_dir, os.path.basename(file_key)) self.s3_client.download_file(self.bucket_name, file_key, local_file_path) return local_dir 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: model_dir = await asyncio.to_thread(self._download_from_s3, model_name) config = AutoConfig.from_pretrained(model_dir) tokenizer = AutoTokenizer.from_pretrained(model_dir, config=config) model = AutoModelForCausalLM.from_pretrained(model_dir, config=config) logging.info(f"Model {model_name} loaded from S3 successfully.") 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) 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 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)