app.py

import os
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, field_validator
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig,
    StoppingCriteria,
    StoppingCriteriaList,
    pipeline
)
import uvicorn
import asyncio
import json
import base64
from huggingface_hub import login
from botocore.exceptions import NoCredentialsError
from functools import lru_cache
from typing import AsyncGenerator


HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")

if HUGGINGFACE_HUB_TOKEN:
    login(token=HUGGINGFACE_HUB_TOKEN,
          add_to_git_credential=False)


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_data = {}  # Global dictionary to store model data

model_load_lock = asyncio.Lock() # Lock to avoid race conditions

@lru_cache(maxsize=None)
async def _load_model_and_tokenizer(model_name):
    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
          
          return {"model":model, "tokenizer":tokenizer}
    except Exception as e:
        raise HTTPException(
            status_code=500, detail=f"Error loading model: {e}"
        )

async def load_model_and_tokenizer(model_name):
    async with model_load_lock:
      if model_name in model_data:
          return model_data[model_name].get("model"), model_data[model_name].get("tokenizer")
      
      model_bundle = await _load_model_and_tokenizer(model_name)
      model_data[model_name] = model_bundle
      return model_bundle.get("model"), model_bundle.get("tokenizer")


@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 load_model_and_tokenizer(model_name)
        device = "cpu" # Force 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,
                eos_token_id = tokenizer.eos_token_id
            )
            if stream:
                return StreamingResponse(
                        stream_json_responses(stream_text(model, tokenizer, input_text,
                                     generation_config, stop_sequences,
                                     device)),
                        media_type="text/plain"
                    )
            else:
                result = await generate_text(model, tokenizer, input_text,
                                     generation_config, stop_sequences,
                                     device)
                return JSONResponse({"text": result, "is_end": True})
        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)}"
        )

class StopOnSequences(StoppingCriteria):
    def __init__(self, stop_sequences, tokenizer):
        self.stop_sequences = stop_sequences
        self.tokenizer = tokenizer
        self.stop_ids = [tokenizer.encode(seq, add_special_tokens=False) for seq in stop_sequences]

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        
        decoded_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)

        for stop_sequence in self.stop_sequences:
             if stop_sequence in decoded_text:
                 return True
        return False


async def stream_text(model, tokenizer, input_text,
                        generation_config, stop_sequences,
                        device) -> AsyncGenerator[dict, None]:
    
    encoded_input = tokenizer(
        input_text, return_tensors="pt",
        truncation=True
    ).to(device)
    
    stop_criteria = StopOnSequences(stop_sequences, tokenizer)
    stopping_criteria = StoppingCriteriaList([stop_criteria])

    output_text = ""
    
    while True:
        
        outputs = await asyncio.to_thread(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,
                                            stopping_criteria=stopping_criteria
                                         )
    
        new_text = tokenizer.decode(
            outputs.sequences[0][len(encoded_input["input_ids"][0]):],
            skip_special_tokens=True
        )

        if not new_text:
            if not stop_criteria(outputs.sequences, None):
                yield {"text": output_text, "is_end": False}
                yield {"text": "", "is_end": True}
            break
        
        output_text += new_text
        
        yield {"text": new_text, "is_end": False}
              
        if stop_criteria(outputs.sequences, None):
            yield {"text": "", "is_end": True}
            break
        
        encoded_input = tokenizer(
            output_text, return_tensors="pt",
            truncation=True
        ).to(device)
        output_text = ""


async def stream_json_responses(generator: AsyncGenerator[dict, None]) -> AsyncGenerator[str, None]:
    async for data in generator:
        yield json.dumps(data) + "\n"


async def generate_text(model, tokenizer, input_text,
                        generation_config, stop_sequences,
                        device):
    encoded_input = tokenizer(
        input_text, return_tensors="pt",
        truncation=True
    ).to(device)
    
    stop_criteria = StopOnSequences(stop_sequences, tokenizer)
    stopping_criteria = StoppingCriteriaList([stop_criteria])
    
    outputs = await asyncio.to_thread(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,
                                        stopping_criteria=stopping_criteria
                                     )
    
    
    generated_text = tokenizer.decode(
        outputs.sequences[0], skip_special_tokens=True
    )
    
    return generated_text

@app.post("/generate-image")
async def generate_image(request: GenerateRequest):
    try:
        validated_body = request
        device = "cpu" # Force CPU
        
        if validated_body.model_name not in model_data:
            config = AutoConfig.from_pretrained(
                    validated_body.model_name, token=HUGGINGFACE_HUB_TOKEN
                )
            model = pipeline(
                "text-to-image", model=validated_body.model_name,
                device=device, config=config
            )
            model_data[validated_body.model_name] = {"model":model}
        else:
            model = model_data[validated_body.model_name]["model"]
        
        image = model(validated_body.input_text)[0]
        
        image_data = list(image.getdata())
        
        return JSONResponse({"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 = "cpu"  # Force CPU

        if validated_body.model_name not in model_data:
             config = AutoConfig.from_pretrained(
                    validated_body.model_name, token=HUGGINGFACE_HUB_TOKEN
                )
            
             audio_generator = pipeline(
                "text-to-speech", model=validated_body.model_name,
                device=device, config=config
            )
             model_data[validated_body.model_name] = {"model":audio_generator}
        else:
            audio_generator = model_data[validated_body.model_name]["model"]

        audio = audio_generator(validated_body.input_text)

        
        audio_bytes = audio["audio"]
        
        audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
        
        return JSONResponse({"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 = "cpu"  # Force CPU
        if validated_body.model_name not in model_data:
            config = AutoConfig.from_pretrained(
                    validated_body.model_name, token=HUGGINGFACE_HUB_TOKEN
                )
            
            video_generator = pipeline(
                "text-to-video", model=validated_body.model_name,
                device=device, config=config
            )
            model_data[validated_body.model_name] = {"model":video_generator}
        else:
            video_generator = model_data[validated_body.model_name]["model"]

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

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

@app.on_event("startup")
async def startup_event():
    # Load models here
    print("Loading models...")
    
    models_to_load = set()
    
    for env_var_key, env_var_value in os.environ.items():
        if env_var_key.startswith("MODEL_NAME_"):
            models_to_load.add(env_var_value)
            
    
    for model_name in models_to_load:
        try:
            await load_model_and_tokenizer(model_name)
            print(f"Model {model_name} loaded")
        except Exception as e:
             print(f"Error loading model {model_name}: {e}")
            

    print("Models loaded.")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)