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)