app.py

import os
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, field_validator
from transformers import (
    AutoConfig,
    pipeline,
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig,
)
import boto3
import uvicorn
import asyncio
from transformers import pipeline
import json
from huggingface_hub import login
import base64
from botocore.exceptions import NoCredentialsError
import re


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 HUGGINGFACE_HUB_TOKEN:
    login(token=HUGGINGFACE_HUB_TOKEN,
          add_to_git_credential=False)

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()

class GenerateRequest(BaseModel):
    model_name: str
    input_text: str = ""
    task_type: str
    temperature: float = 1.0
    max_new_tokens: int = 3
    stream: bool = True
    top_p: float = 1.0
    top_k: int = 50    
    repetition_penalty: float = 1.0
    num_return_sequences: int = 1
    do_sample: bool = True
    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

model_cache = {}

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}/" \
               f"{model_name.replace('/', '-')}"

    async def load_model_and_tokenizer(self, model_name):
       if model_name in model_cache:
          return model_cache[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
            )

            if tokenizer.eos_token_id is not None and \
               tokenizer.pad_token_id is None:
                tokenizer.pad_token_id = config.pad_token_id \
                                        or tokenizer.eos_token_id
            model_cache[model_name] = (model, tokenizer)
            return model, tokenizer
       except (EnvironmentError, NoCredentialsError):
            try:
                config = AutoConfig.from_pretrained(
                    model_name, token=HUGGINGFACE_HUB_TOKEN
                )
                tokenizer = AutoTokenizer.from_pretrained(
                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                )
                
                model = AutoModelForCausalLM.from_pretrained(
                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                )


                if tokenizer.eos_token_id is not None and \
                   tokenizer.pad_token_id is None:
                    tokenizer.pad_token_id = config.pad_token_id \
                                            or tokenizer.eos_token_id

                
                model.save_pretrained(s3_uri)
                tokenizer.save_pretrained(s3_uri)
                
                
                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
                )
                model_cache[model_name] = (model, tokenizer)
                return model, tokenizer
            except Exception as e:
                raise HTTPException(
                    status_code=500, detail=f"Error loading model: {e}"
                )

model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)

@app.post("/generate")
async def generate(request: GenerateRequest):
    try:
        model_name = request.model_name
        input_text = request.input_text
        task_type = request.task_type
        temperature = request.temperature
        max_new_tokens = request.max_new_tokens
        stream = request.stream
        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

        model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)
        
        if "text-to-text" == task_type:
            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,
            )

            return StreamingResponse(
                stream_text(model, tokenizer, input_text,
                             generation_config, stop_sequences,
                             device),
                media_type="text/plain"
            )
        else:
            return HTTPException(status_code=400, detail="Task type not text-to-text")

    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Internal server error: {str(e)}"
        )


async def stream_text(model, tokenizer, input_text,
                        generation_config, stop_sequences,
                        device):

    encoded_input = tokenizer(
        input_text, return_tensors="pt",
        truncation=True
    ).to(device)
    
    stop_regex = re.compile(r'[\.\?\!\n]+')
    
    def find_stop(output_text, stop_sequences):
        for seq in stop_sequences:
           if seq in output_text:
              last_index = output_text.rfind(seq)
              return last_index + len(seq)
        
        match = stop_regex.search(output_text)
        if match:
           return match.end()
        
        return -1
    
    
    output_text = ""
    while True:
        outputs = model.generate(
            **encoded_input,
            do_sample=generation_config.do_sample,
            max_new_tokens=generation_config.max_new_tokens,
            temperature=generation_config.temperature,
            top_p=generation_config.top_p,
            top_k=generation_config.top_k,
            repetition_penalty=generation_config.repetition_penalty,
            num_return_sequences=generation_config.num_return_sequences,
            output_scores=True,
            return_dict_in_generate=True,
        )

        new_text = tokenizer.decode(
            outputs.sequences[0][len(encoded_input["input_ids"][0]):],
            skip_special_tokens=True
        )
        
        output_text += new_text

        stop_index = find_stop(output_text, stop_sequences)

        if stop_index != -1:
            final_output = output_text[:stop_index]
            
            for text in final_output.split():
                 yield json.dumps({"text": text, "is_end": False}) + "\n"
            yield json.dumps({"text": "", "is_end": True}) + "\n"
            break
        else:
            for text in new_text.split():
              yield json.dumps({"text": text, "is_end": False}) + "\n"
        
        if len(new_text) == 0:
            for text in output_text.split():
                 yield json.dumps({"text": text, "is_end": False}) + "\n"
            yield json.dumps({"text": "", "is_end": True}) + "\n"
            break

        encoded_input = tokenizer(
            output_text, return_tensors="pt",
            truncation=True
        ).to(device)
        output_text = ""



@app.post("/generate-image")
async def generate_image(request: GenerateRequest):
    try:
        validated_body = request
        device = "cuda" if torch.cuda.is_available() else "cpu"
        
        if validated_body.model_name not in model_cache:
            model = pipeline(
                "text-to-image", model=validated_body.model_name,
                device=device
            )
            model_cache[validated_body.model_name] = model
        else:
            model = model_cache[validated_body.model_name]
        
        image = model(validated_body.input_text)[0]
        
        image_data = list(image.getdata())
        
        return json.dumps({"image_data": image_data, "is_end": True})

    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:
        validated_body = request
        device = "cuda" if torch.cuda.is_available() else "cpu"

        if validated_body.model_name not in model_cache:
            audio_generator = pipeline(
                "text-to-speech", model=validated_body.model_name,
                device=device
            )
            model_cache[validated_body.model_name] = audio_generator
        else:
            audio_generator = model_cache[validated_body.model_name]

        audio = audio_generator(validated_body.input_text)

        
        audio_bytes = audio["audio"]
        
        audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
        
        return json.dumps({"audio": audio_base64, "is_end": True})

    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:
        validated_body = request
        device = "cuda" if torch.cuda.is_available() else "cpu"
        if validated_body.model_name not in model_cache:
            video_generator = pipeline(
            "text-to-video", model=validated_body.model_name,
            device=device
        )
            model_cache[validated_body.model_name] = video_generator
        else:
            video_generator = model_cache[validated_body.model_name]

        video = video_generator(validated_body.input_text)
        
        
        video_base64 = base64.b64encode(video).decode('utf-8')
        
        return json.dumps({"video": video_base64, "is_end": True})

    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)