import os from fastapi import FastAPI, HTTPException, Depends from fastapi.responses import JSONResponse from pydantic import BaseModel, field_validator from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList, pipeline import boto3 import uvicorn import soundfile as sf import imageio from typing import Dict 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") if not all([AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION, S3_BUCKET_NAME]): raise ValueError("Missing one or more AWS environment variables.") 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) app = FastAPI() SPECIAL_TOKENS = { "bos_token": "<|startoftext|>", "eos_token": "<|endoftext|>", "pad_token": "[PAD]", "unk_token": "[UNK]", } class GenerateRequest(BaseModel): model_name: str input_text: str = "" task_type: str temperature: float = 1.0 max_new_tokens: int = 10 top_p: float = 1.0 top_k: int = 50 repetition_penalty: float = 1.1 num_return_sequences: int = 1 do_sample: bool = True stop_sequences: list[str] = [] no_repeat_ngram_size: int = 2 continuation_id: str = None @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 @field_validator("max_new_tokens") def max_new_tokens_must_be_within_limit(cls, v): if v > 500: raise ValueError("max_new_tokens cannot be greater than 500.") 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: config = AutoConfig.from_pretrained(s3_uri, local_files_only=False) model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=False) tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config, local_files_only=False) tokenizer.add_special_tokens(SPECIAL_TOKENS) model.resize_token_embeddings(len(tokenizer)) if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id return model, tokenizer except Exception as e: raise HTTPException(status_code=500, detail=f"Error loading model from S3: {e}") model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client) active_generations: Dict[str, Dict] = {} async def get_model_and_tokenizer(model_name: str): try: return await model_loader.load_model_and_tokenizer(model_name) except Exception as e: raise HTTPException(status_code=500, detail=f"Error loading model: {e}") @app.post("/generate") async def generate(request: GenerateRequest, model_resources: tuple = Depends(get_model_and_tokenizer)): model, tokenizer = model_resources try: model_name = request.model_name input_text = request.input_text temperature = request.temperature max_new_tokens = request.max_new_tokens top_p = request.top_p top_k = request.top_k repetition_penalty = request.repetition_penalty num_return_sequences = request.num_return_sequences do_sample = request.do_sample stop_sequences = request.stop_sequences no_repeat_ngram_size = request.no_repeat_ngram_size continuation_id = request.continuation_id if continuation_id: if continuation_id not in active_generations: raise HTTPException(status_code=404, detail="Continuation ID not found.") previous_output = active_generations[continuation_id]["output"] input_text = previous_output generation_config = GenerationConfig( temperature=temperature, max_new_tokens=max_new_tokens, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, do_sample=do_sample, num_return_sequences=num_return_sequences, no_repeat_ngram_size=no_repeat_ngram_size, pad_token_id=tokenizer.pad_token_id ) generated_text = generate_text_internal(model, tokenizer, input_text, generation_config, stop_sequences) if not continuation_id: continuation_id = os.urandom(16).hex() active_generations[continuation_id] = {"model_name": model_name, "output": generated_text} else: active_generations[continuation_id]["output"] = generated_text return JSONResponse({"text": generated_text, "continuation_id": continuation_id}) except HTTPException as http_err: raise http_err except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") def generate_text_internal(model, tokenizer, input_text, generation_config, stop_sequences): max_model_length = model.config.max_position_embeddings encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True) stopping_criteria = StoppingCriteriaList() class CustomStoppingCriteria(StoppingCriteriaList): def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: decoded_output = tokenizer.decode(input_ids[0], skip_special_tokens=True) for stop in stop_sequences: if decoded_output.endswith(stop): return True return False stopping_criteria.append(CustomStoppingCriteria()) outputs = model.generate( encoded_input.input_ids, generation_config=generation_config, stopping_criteria=stopping_criteria, pad_token_id=generation_config.pad_token_id ) generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) return generated_text async def load_pipeline_from_s3(task, model_name): s3_uri = f"s3://{S3_BUCKET_NAME}/{model_name.replace('/', '-')}" try: return pipeline(task, model=s3_uri) except Exception as e: raise HTTPException(status_code=500, detail=f"Error loading {task} model from S3: {e}") @app.post("/generate-image") async def generate_image(request: GenerateRequest): try: if request.task_type != "text-to-image": raise HTTPException(status_code=400, detail="Invalid task_type for this endpoint.") image_generator = await load_pipeline_from_s3("text-to-image", request.model_name) image = image_generator(request.input_text)[0] continuation_id = os.urandom(16).hex() active_generations[continuation_id] = {"model_name": request.model_name, "output": "Image generated successfully"} return JSONResponse({"url": "Image generated successfully", "continuation_id": continuation_id}) except HTTPException as http_err: raise http_err except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") @app.post("/generate-text-to-speech") async def generate_text_to_speech(request: GenerateRequest): try: if request.task_type != "text-to-speech": raise HTTPException(status_code=400, detail="Invalid task_type for this endpoint.") tts_pipeline = await load_pipeline_from_s3("text-to-speech", request.model_name) output = tts_pipeline(request.input_text) continuation_id = os.urandom(16).hex() active_generations[continuation_id] = {"model_name": request.model_name, "output": "Audio generated successfully"} return JSONResponse({"url": "Audio generated successfully", "continuation_id": continuation_id}) except HTTPException as http_err: raise http_err except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") @app.post("/generate-video") async def generate_video(request: GenerateRequest): try: if request.task_type != "text-to-video": raise HTTPException(status_code=400, detail="Invalid task_type for this endpoint.") video_pipeline = await load_pipeline_from_s3("text-to-video", request.model_name) output = video_pipeline(request.input_text) continuation_id = os.urandom(16).hex() active_generations[continuation_id] = {"model_name": request.model_name, "output": "Video generated successfully"} return JSONResponse({"url": "Video generated successfully", "continuation_id": continuation_id}) except HTTPException as http_err: raise http_err except Exception as e: raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=7860)