app.py

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)