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)