Update app.py

app.py

@@ -1,197 +1,155 @@
 import os
+import json
+import threading
 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, field_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
+from google.cloud import storage
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from pydantic import BaseModel
+from fastapi import FastAPI, HTTPException
+import requests
 import uvicorn
+from dotenv import load_dotenv
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
+load_dotenv()
 
-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")
+API_KEY = os.getenv("API_KEY")
+GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
+GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
+HF_API_TOKEN = os.getenv("HF_API_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] = []
-
-    model_config = {"protected_namespaces": ()}
-
-    @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
-
-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, config=config)
-
-            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
-
-            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:
-                config = AutoConfig.from_pretrained(model_name)
-                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
-                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
-
-                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}")
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-app = FastAPI()
+credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
+storage_client = storage.Client.from_service_account_info(credentials_info)
+bucket = storage_client.bucket(GCS_BUCKET_NAME)
 
-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 = FastAPI()
 
-@app.post("/generate")
-async def generate(request: Request, body: GenerateRequest):
+class DownloadModelRequest(BaseModel):
+    model_name: str
+    pipeline_task: str
+    input_text: str
+
+class GCSHandler:
+    def __init__(self, bucket_name):
+        self.bucket = storage_client.bucket(bucket_name)
+
+    def file_exists(self, blob_name):
+        return self.bucket.blob(blob_name).exists()
+
+    def create_folder_if_not_exists(self, folder_name):
+        if not self.file_exists(folder_name):
+            self.bucket.blob(folder_name + "/").upload_from_string("")
+
+    def upload_file(self, blob_name, file_stream):
+        self.create_folder_if_not_exists(os.path.dirname(blob_name))
+        blob = self.bucket.blob(blob_name)
+        blob.upload_from_file(file_stream)
+
+    def download_file(self, blob_name):
+        blob = self.bucket.blob(blob_name)
+        if not blob.exists():
+            raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
+        return blob.open("rb")
+
+    def generate_signed_url(self, blob_name, expiration=3600):
+        blob = self.bucket.blob(blob_name)
+        return blob.generate_signed_url(expiration=expiration)
+
+def download_model_from_huggingface(model_name):
+    url = f"https://huggingface.co/{model_name}/tree/main"
+    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
+    response = requests.get(url, headers=headers)
+    if response.status_code == 200:
+        model_files = [
+            "pytorch_model.bin",
+            "config.json",
+            "tokenizer.json",
+            "model.safetensors",
+        ]
+        for file_name in model_files:
+            file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
+            file_content = requests.get(file_url).content
+            blob_name = f"models/{model_name}/{file_name}"
+            bucket.blob(blob_name).upload_from_string(file_content)
+    else:
+        raise HTTPException(status_code=404, detail="Error accessing Hugging Face model files.")
+
+def download_and_verify_model(model_name):
+    model_files = [
+        "pytorch_model.bin",
+        "config.json",
+        "tokenizer.json",
+        "model.safetensors",
+    ]
+    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
+    if not all(gcs_handler.file_exists(f"models/{model_name}/{file}") for file in model_files):
+        download_model_from_huggingface(model_name)
+
+def load_model_from_gcs(model_name):
+    model_files = [
+        "pytorch_model.bin",
+        "config.json",
+        "tokenizer.json",
+        "model.safetensors",
+    ]
+    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
+    model_files_streams = {
+        file: gcs_handler.download_file(f"models/{model_name}/{file}")
+        for file in model_files if gcs_handler.file_exists(f"models/{model_name}/{file}")
+    }
+    model_stream = model_files_streams.get("pytorch_model.bin") or model_files_streams.get("model.safetensors")
+    tokenizer_stream = model_files_streams.get("tokenizer.json")
+    config_stream = model_files_streams.get("config.json")
+    model = AutoModelForCausalLM.from_pretrained(model_stream, config=config_stream)
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
+    return model, tokenizer
+
+def load_model(model_name):
+    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
     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
-            )
-
-            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)
-
-                    stopping_criteria = StoppingCriteriaList(
-                        [lambda _, outputs: tokenizer.decode(outputs[0][-1], skip_special_tokens=True) in validated_body.stop_sequences] if validated_body.stop_sequences else []
-                    )
-
-                    output = model.generate(**encoded_input, generation_config=generation_config, stopping_criteria=stopping_criteria)
-                    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.")
-
+        return load_model_from_gcs(model_name)
+    except HTTPException:
+        download_and_verify_model(model_name)
+        return load_model_from_gcs(model_name)
+
+@app.on_event("startup")
+async def startup():
+    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
+    blobs = list(bucket.list_blobs(prefix="models/"))
+    model_names = set(blob.name.split("/")[1] for blob in blobs)
+    def download_model_thread(model_name):
+        try:
+            download_and_verify_model(model_name)
+        except Exception as e:
+            logger.error(f"Error downloading model '{model_name}': {e}")
+    threads = [threading.Thread(target=download_model_thread, args=(model_name,)) for model_name in model_names]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+@app.post("/predict/")
+async def predict(request: DownloadModelRequest):
+    model_name = request.model_name
+    pipeline_task = request.pipeline_task
+    input_text = request.input_text
+    model, tokenizer = load_model(model_name)
+    pipe = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
+    result = pipe(input_text)
+    return {"result": result}
+
+def download_all_models_in_background():
+    models_url = "https://huggingface.co/api/models"
+    response = requests.get(models_url)
+    if response.status_code == 200:
+        models = response.json()
+        for model in models:
+            download_model_from_huggingface(model["id"])
+
+def run_in_background():
+    threading.Thread(target=download_all_models_in_background, daemon=True).start()
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)