app.py

import os
import torch
from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Depends, BackgroundTasks, Request, Query, APIRouter, Path, Body, status, Response, Header
from fastapi.responses import StreamingResponse, JSONResponse, FileResponse, HTMLResponse, PlainTextResponse, RedirectResponse
from pydantic import BaseModel, validator, Field, root_validator, EmailStr, constr, ValidationError
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig,
    StoppingCriteriaList,
    pipeline,
    AutoProcessor,
    AutoModelForImageClassification,
    AutoModelForSeq2SeqLM,
    AutoModelForQuestionAnswering,
    AutoModelForSpeechSeq2Seq,
    AutoModelForImageSegmentation,
    AutoFeatureExtractor,
    AutoModelForTokenClassification,
    AutoModelForMaskedLM,
    AutoModelForObjectDetection,
    AutoModelForSeq2SeqLM
)
from io import BytesIO
import boto3
from botocore.exceptions import NoCredentialsError, ClientError
from huggingface_hub import snapshot_download
import asyncio
import tempfile
import hashlib
from PIL import Image
import base64
from typing import Optional, List, Union, Dict, Any
import uuid
import subprocess
import json
from starlette.middleware.cors import CORSMiddleware
import numpy as np
from typing import Dict, Any
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.middleware.gzip import GZipMiddleware
from transformers import AutoImageProcessor, pipeline
from fastapi.security import APIKeyHeader, OAuth2PasswordBearer, OAuth2PasswordRequestForm
from fastapi.security.api_key import APIKeyCookie
from fastapi import Depends, Security, status, APIRouter, UploadFile, File, Request
from fastapi.security import APIKeyHeader, OAuth2PasswordRequestForm
from passlib.context import CryptContext
from jose import JWTError, jwt
from datetime import datetime, timedelta
from starlette.requests import Request
import logging
from pydantic import EmailStr, constr, ValidationError
from database import insert_user, get_user, delete_user, update_user, create_db_and_table
from starlette.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.types import ASGIApp
import uvicorn
from starlette.responses import StreamingResponse
import logging
from pydantic import EmailStr, constr, ValidationError
from database import insert_user, get_user, delete_user, update_user, create_db_and_table, get_all_users
from starlette.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.types import ASGIApp
import uvicorn
from starlette.responses import StreamingResponse
import logging
from fastapi.exceptions import RequestValidationError
from fastapi import Request, status, Depends
from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer
from jose import JWTError, jwt
from passlib.context import CryptContext
from datetime import datetime, timedelta
from pydantic import BaseModel, field_validator, model_validator, Field, EmailStr, constr, ValidationError
from typing import Optional, List, Union

#setting up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(filename)s - %(lineno)d - %(message)s')
logger = logging.getLogger(__name__)

#JWT Settings
SECRET_KEY = os.getenv("SECRET_KEY")
if not SECRET_KEY:
    raise ValueError("SECRET_KEY must be set.")
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

#Password Hashing
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

#Database connection - replace with your database setup
#Example using SQLite
import sqlite3
conn = sqlite3.connect('users.db')
cursor = conn.cursor()

#OAuth2
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

#API Key
API_KEY = os.getenv("API_KEY")
api_key_header = APIKeyHeader(name="X-API-Key")

#Configuration
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")
TEMP_DIR = "/tmp"
STATIC_DIR = "static"
TEMPLATES = Jinja2Templates(directory="templates")

app = FastAPI()
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.add_middleware(GZipMiddleware)

origins = ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class User(BaseModel):
    username: constr(min_length=3, max_length=50)
    email: EmailStr
    password: constr(min_length=8)

class GenerateRequest(BaseModel):
    model_name: str
    input_text: Optional[str] = Field(None, description="Input text for generation.")
    task_type: str = Field(..., description="Type of generation task (text, image, audio, video, classification, translation, question-answering, speech-to-text, text-to-speech, image-segmentation, feature-extraction, token-classification, fill-mask, image-inpainting, image-super-resolution, object-detection, image-captioning, audio-transcription, summarization).")
    temperature: float = 1.0
    max_new_tokens: int = 200
    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
    chunk_delay: float = 0.0
    stop_sequences: List[str] = []
    image_file: Optional[UploadFile] = None
    source_language: Optional[str] = None
    target_language: Optional[str] = None
    context: Optional[str] = None
    audio_file: Optional[UploadFile] = None
    raw_input: Optional[Union[str, bytes]] = None  # for feature extraction
    masked_text: Optional[str] = None  # for fill-mask
    mask_image: Optional[UploadFile] = None  # for image inpainting
    low_res_image: Optional[UploadFile] = None  # for image super-resolution

    @field_validator('task_type')
    def validate_task_type(cls, value):
        allowed_types = ["text", "image", "audio", "video", "classification", "translation", "question-answering", "speech-to-text", "text-to-speech", "image-segmentation", "feature-extraction", "token-classification", "fill-mask", "image-inpainting", "image-super-resolution", "object-detection", "image-captioning", "audio-transcription", "summarization"]
        if value not in allowed_types:
            raise ValueError(f"Invalid task_type. Allowed types are: {allowed_types}")
        return value

    @model_validator(mode='after')
    def check_input(cls, values):
        task_type = values.get("task_type")
        if task_type == "text" and values.get("input_text") is None:
            raise ValueError("input_text is required for text generation.")
        elif task_type == "speech-to-text" and values.get("audio_file") is None:
            raise ValueError("audio_file is required for speech-to-text.")
        elif task_type == "classification" and values.get("image_file") is None:
            raise ValueError("image_file is required for image classification.")
        elif task_type == "image-segmentation" and values.get("image_file") is None:
            raise ValueError("image_file is required for image segmentation.")
        elif task_type == "feature-extraction" and values.get("raw_input") is None:
            raise ValueError("raw_input is required for feature extraction.")
        elif task_type == "fill-mask" and values.get("masked_text") is None:
            raise ValueError("masked_text is required for fill-mask.")
        elif task_type == "image-inpainting" and (values.get("image_file") is None or values.get("mask_image") is None):
            raise ValueError("image_file and mask_image are required for image inpainting.")
        elif task_type == "image-super-resolution" and values.get("low_res_image") is None:
            raise ValueError("low_res_image is required for image super-resolution.")
        return values

class S3ModelLoader:
    def __init__(self, bucket_name, aws_access_key_id, aws_secret_access_key, aws_region):
        self.bucket_name = bucket_name
        self.s3 = boto3.client(
            's3',
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key,
            region_name=aws_region
        )

    def _get_s3_uri(self, model_name):
        return f"{self.bucket_name}/{model_name.replace('/', '-')}"

    def load_model_and_tokenizer(self, model_name, task_type):
        s3_uri = self._get_s3_uri(model_name)
        try:
            self.s3.head_object(Bucket=self.bucket_name, Key=f'{s3_uri}/config.json')
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                with tempfile.TemporaryDirectory() as tmpdir:
                    model_path = snapshot_download(model_name, token=HUGGINGFACE_HUB_TOKEN, cache_dir=tmpdir)
                    self._upload_model_to_s3(model_path, s3_uri)
            else:
                raise HTTPException(status_code=500, detail=f"Error accessing S3: {e}")
        return self._load_from_s3(s3_uri, task_type)

    def _upload_model_to_s3(self, model_path, s3_uri):
        for root, _, files in os.walk(model_path):
            for file in files:
                local_path = os.path.join(root, file)
                s3_path = os.path.join(s3_uri, os.path.relpath(local_path, model_path))
                self.s3.upload_file(local_path, self.bucket_name, s3_path)

    def _load_from_s3(self, s3_uri, task_type):
        with tempfile.TemporaryDirectory() as tmpdir:
            model_path = os.path.join(tmpdir, s3_uri)
            os.makedirs(model_path, exist_ok=True)
            self.s3.download_file(self.bucket_name, f"{s3_uri}/config.json", os.path.join(model_path, "config.json"))
            if task_type == "text":
                model = AutoModelForCausalLM.from_pretrained(model_path, load_in_8bit=True)
                tokenizer = AutoTokenizer.from_pretrained(model_path)
                if tokenizer.eos_token_id is None:
                    tokenizer.eos_token_id = tokenizer.pad_token_id
                return {"model": model, "tokenizer": tokenizer, "pad_token_id": tokenizer.pad_token_id, "eos_token_id": tokenizer.eos_token_id}
            elif task_type in ["image", "audio", "video"]:
                processor = AutoProcessor.from_pretrained(model_path)
                pipeline_function = pipeline(task_type, model=model_path, device=0 if torch.cuda.is_available() else -1, processor=processor)
                return {"pipeline": pipeline_function}
            elif task_type == "classification":
                model = AutoModelForImageClassification.from_pretrained(model_path)
                processor = AutoProcessor.from_pretrained(model_path)
                return {"model": model, "processor": processor}
            elif task_type == "translation":
                model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
                tokenizer = AutoTokenizer.from_pretrained(model_path)
                return {"model": model, "tokenizer": tokenizer}
            elif task_type == "question-answering":
                model = AutoModelForQuestionAnswering.from_pretrained(model_path)
                tokenizer = AutoTokenizer.from_pretrained(model_path)
                return {"model": model, "tokenizer": tokenizer}
            elif task_type == "speech-to-text":
                model = pipeline("automatic-speech-recognition", model=model_path, device=0 if torch.cuda.is_available() else -1)
                return {"pipeline": model}
            elif task_type == "text-to-speech":
                model = pipeline("text-to-speech", model=model_path, device=0 if torch.cuda.is_available() else -1)
                return {"pipeline": model}
            elif task_type == "image-segmentation":
                model = pipeline("image-segmentation", model=model_path, device=0 if torch.cuda.is_available() else -1)
                return {"pipeline": model}
            elif task_type == "feature-extraction":
                feature_extractor = AutoFeatureExtractor.from_pretrained(model_path)
                return {"feature_extractor": feature_extractor}
            elif task_type == "token-classification":
                model = AutoModelForTokenClassification.from_pretrained(model_path)
                tokenizer = AutoTokenizer.from_pretrained(model_path)
                return {"model": model, "tokenizer": tokenizer}
            elif task_type == "fill-mask":
                model = AutoModelForMaskedLM.from_pretrained(model_path)
                tokenizer = AutoTokenizer.from_pretrained(model_path)
                return {"model": model, "tokenizer": tokenizer}
            elif task_type == "image-inpainting":
                model = pipeline("image-inpainting", model=model_path, device=0 if torch.cuda.is_available() else -1)
                return {"pipeline": model}
            elif task_type == "image-super-resolution":
                model = pipeline("image-super-resolution", model=model_path, device=0 if torch.cuda.is_available() else -1)
                return {"pipeline": model}
            elif task_type == "object-detection":
                model = pipeline("object-detection", model=model_path, device=0 if torch.cuda.is_available() else -1)
                image_processor = AutoImageProcessor.from_pretrained(model_path)
                return {"pipeline": model, "image_processor": image_processor}
            elif task_type == "image-captioning":
                model = pipeline("image-captioning", model=model_path, device=0 if torch.cuda.is_available() else -1)
                return {"pipeline": model}
            elif task_type == "audio-transcription":
                model = pipeline("automatic-speech-recognition", model=model_path, device=0 if torch.cuda.is_available() else -1)
                return {"pipeline": model}
            elif task_type == "summarization":
                model = pipeline("summarization", model=model_path, device=0 if torch.cuda.is_available() else -1)
                tokenizer = AutoTokenizer.from_pretrained(model_path)
                return {"model": model, "tokenizer": tokenizer}
            else:
                raise ValueError("Unsupported task type")

async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay):
    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True).to(device)
    input_length = encoded_input["input_ids"].shape[1]
    max_length = model.config.max_length
    remaining_tokens = max_length - input_length
    if remaining_tokens <= 0:
        yield ""
    generation_config.max_new_tokens = min(remaining_tokens, generation_config.max_new_tokens)
    def stop_criteria(input_ids, scores):
        decoded_output = tokenizer.decode(input_ids[0][-1], skip_special_tokens=True)
        return decoded_output in stop_sequences
    stopping_criteria = StoppingCriteriaList([stop_criteria])
    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,
        stopping_criteria=stopping_criteria,
        output_scores=True,
        return_dict_in_generate=True
    )
    for output in outputs.sequences:
        for token_id in output:
            token = tokenizer.decode(token_id, skip_special_tokens=True)
            yield token


model_loader = S3ModelLoader(S3_BUCKET_NAME, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION)

def get_model_data(request: GenerateRequest):
    return model_loader.load_model_and_tokenizer(request.model_name, request.task_type)

async def verify_api_key(api_key: str = Depends(api_key_header)):
    if api_key != API_KEY:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API Key")


@app.post("/generate", dependencies=[Depends(verify_api_key)])
async def generate(request: GenerateRequest, background_tasks: BackgroundTasks, model_data = Depends(get_model_data)):
    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        if request.task_type == "text":
            model = model_data["model"].to(device)
            tokenizer = model_data["tokenizer"]
            generation_config = GenerationConfig(
                temperature=request.temperature,
                max_new_tokens=request.max_new_tokens,
                top_p=request.top_p,
                top_k=request.top_k,
                repetition_penalty=request.repetition_penalty,
                do_sample=request.do_sample,
                num_return_sequences=request.num_return_sequences,
            )
            async def stream_with_tokens():
                async for token in stream_text(model, tokenizer, request.input_text, generation_config, request.stop_sequences, device, request.chunk_delay):
                    yield f"Token: {token}\n"
            return StreamingResponse(stream_with_tokens(), media_type="text/plain")
        elif request.task_type in ["image", "audio", "video"]:
            pipeline = model_data["pipeline"]
            result = pipeline(request.input_text)
            if request.task_type == "image":
                image = result[0]
                img_byte_arr = BytesIO()
                image.save(img_byte_arr, format="PNG")
                img_byte_arr.seek(0)
                return StreamingResponse(img_byte_arr, media_type="image/png")
            elif request.task_type == "audio":
                audio = result[0]
                audio_byte_arr = BytesIO()
                audio.save(audio_byte_arr, format="wav")
                audio_byte_arr.seek(0)
                return StreamingResponse(audio_byte_arr, media_type="audio/wav")
            elif request.task_type == "video":
                video = result[0]
                video_byte_arr = BytesIO()
                video.save(video_byte_arr, format="mp4")
                video_byte_arr.seek(0)
                return StreamingResponse(video_byte_arr, media_type="video/mp4")
        elif request.task_type == "classification":
            if request.image_file is None:
                raise HTTPException(status_code=400, detail="Image file is required for classification.")
            contents = await request.image_file.read()
            image = Image.open(BytesIO(contents)).convert("RGB")
            model = model_data["model"].to(device)
            processor = model_data["processor"]
            inputs = processor(images=image, return_tensors="pt").to(device)
            with torch.no_grad():
                outputs = model(**inputs)
            predicted_class_idx = outputs.logits.argmax().item()
            predicted_class = model.config.id2label[predicted_class_idx]
            return JSONResponse({"predicted_class": predicted_class})
        elif request.task_type == "translation":
            if request.source_language is None or request.target_language is None:
                raise HTTPException(status_code=400, detail="Source and target languages are required for translation.")
            model = model_data["model"].to(device)
            tokenizer = model_data["tokenizer"]
            inputs = tokenizer(request.input_text, return_tensors="pt").to(device)
            with torch.no_grad():
                outputs = model.generate(**inputs)
            translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
            return JSONResponse({"translation": translation})
        elif request.task_type == "question-answering":
            if request.context is None:
                raise HTTPException(status_code=400, detail="Context is required for question answering.")
            model = model_data["model"].to(device)
            tokenizer = model_data["tokenizer"]
            inputs = tokenizer(question=request.input_text, context=request.context, return_tensors="pt").to(device)
            with torch.no_grad():
                outputs = model(**inputs)
            answer_start = torch.argmax(outputs.start_logits)
            answer_end = torch.argmax(outputs.end_logits) + 1
            answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))
            return JSONResponse({"answer": answer})
        elif request.task_type == "speech-to-text":
            if request.audio_file is None:
                raise HTTPException(status_code=400, detail="Audio file is required for speech-to-text.")
            contents = await request.audio_file.read()
            pipeline = model_data["pipeline"]
            try:
                transcription = pipeline(contents, sampling_rate=16000)[0]["text"]  # Assuming 16kHz sampling rate
                return JSONResponse({"transcription": transcription})
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Error during speech-to-text: {str(e)}")

        elif request.task_type == "text-to-speech":
            if not request.input_text:
                raise HTTPException(status_code=400, detail="Input text is required for text-to-speech.")
            pipeline = model_data["pipeline"]
            try:
                audio = pipeline(request.input_text)[0]
                file_path = os.path.join(TEMP_DIR, f"{uuid.uuid4()}.wav")
                audio.save(file_path)
                background_tasks.add_task(os.remove, file_path)
                return FileResponse(file_path, media_type="audio/wav")
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Error during text-to-speech: {str(e)}")

        elif request.task_type == "image-segmentation":
            if request.image_file is None:
                raise HTTPException(status_code=400, detail="Image file is required for image segmentation.")
            contents = await request.image_file.read()
            image = Image.open(BytesIO(contents)).convert("RGB")
            pipeline = model_data["pipeline"]
            result = pipeline(image)
            mask = result[0]['mask']
            mask_byte_arr = BytesIO()
            mask.save(mask_byte_arr, format="PNG")
            mask_byte_arr.seek(0)
            return StreamingResponse(mask_byte_arr, media_type="image/png")
        elif request.task_type == "feature-extraction":
            if request.raw_input is None:
                raise HTTPException(status_code=400, detail="raw_input is required for feature extraction.")
            feature_extractor = model_data["feature_extractor"]
            try:
                if isinstance(request.raw_input, str):
                    inputs = feature_extractor(text=request.raw_input, return_tensors="pt")
                elif isinstance(request.raw_input, bytes):
                    image = Image.open(BytesIO(request.raw_input)).convert("RGB")
                    inputs = feature_extractor(images=image, return_tensors="pt")
                else:
                    raise ValueError("Unsupported raw_input type.")
                features = inputs.pixel_values  # Adjust according to your feature extractor
                return JSONResponse({"features": features.tolist()})
            except Exception as fe:
                raise HTTPException(status_code=400, detail=f"Error during feature extraction: {fe}")
        elif request.task_type == "token-classification":
            if request.input_text is None:
                raise HTTPException(status_code=400, detail="Input text is required for token classification.")
            model = model_data["model"].to(device)
            tokenizer = model_data["tokenizer"]
            inputs = tokenizer(request.input_text, return_tensors="pt", padding=True, truncation=True)
            with torch.no_grad():
                outputs = model(**inputs)
            predictions = outputs.logits.argmax(dim=-1)
            predicted_labels = [model.config.id2label[label_id] for label_id in predictions[0].tolist()]
            return JSONResponse({"predicted_labels": predicted_labels})
        elif request.task_type == "fill-mask":
            if request.masked_text is None:
                raise HTTPException(status_code=400, detail="masked_text is required for fill-mask.")
            model = model_data["model"].to(device)
            tokenizer = model_data["tokenizer"]
            inputs = tokenizer(request.masked_text, return_tensors="pt")
            with torch.no_grad():
                outputs = model(**inputs)
            logits = outputs.logits
            masked_index = torch.where(inputs.input_ids == tokenizer.mask_token_id)[1]
            predicted_token_id = torch.argmax(logits[0, masked_index])
            predicted_token = tokenizer.decode(predicted_token_id)
            return JSONResponse({"predicted_token": predicted_token})
        elif request.task_type == "image-inpainting":
            if request.image_file is None or request.mask_image is None:
                raise HTTPException(status_code=400, detail="image_file and mask_image are required for image inpainting.")
            image_contents = await request.image_file.read()
            mask_contents = await request.mask_image.read()
            image = Image.open(BytesIO(image_contents)).convert("RGB")
            mask = Image.open(BytesIO(mask_contents)).convert("L")  # Assuming mask is grayscale
            pipeline = model_data["pipeline"]
            result = pipeline(image, mask)
            inpainted_image = result[0]
            img_byte_arr = BytesIO()
            inpainted_image.save(img_byte_arr, format="PNG")
            img_byte_arr.seek(0)
            return StreamingResponse(img_byte_arr, media_type="image/png")
        elif request.task_type == "image-super-resolution":
            if request.low_res_image is None:
                raise HTTPException(status_code=400, detail="low_res_image is required for image super-resolution.")
            contents = await request.low_res_image.read()
            image = Image.open(BytesIO(contents)).convert("RGB")
            pipeline = model_data["pipeline"]
            result = pipeline(image)
            upscaled_image = result[0]
            img_byte_arr = BytesIO()
            upscaled_image.save(img_byte_arr, format="PNG")
            img_byte_arr.seek(0)
            return StreamingResponse(img_byte_arr, media_type="image/png")
        elif request.task_type == "object-detection":
            if request.image_file is None:
                raise HTTPException(status_code=400, detail="Image file is required for object detection.")
            contents = await request.image_file.read()
            image = Image.open(BytesIO(contents)).convert("RGB")
            pipeline = model_data["pipeline"]
            image_processor = model_data["image_processor"]
            inputs = image_processor(images=image, return_tensors="pt")
            with torch.no_grad():
                outputs = pipeline(image)
            detections = outputs
            return JSONResponse({"detections": detections})
        elif request.task_type == "image-captioning":
            if request.image_file is None:
                raise HTTPException(status_code=400, detail="Image file is required for image captioning.")
            contents = await request.image_file.read()
            image = Image.open(BytesIO(contents)).convert("RGB")
            pipeline = model_data["pipeline"]
            caption = pipeline(image)[0]['generated_text']
            return JSONResponse({"caption": caption})
        elif request.task_type == "audio-transcription":
            if request.audio_file is None:
                raise HTTPException(status_code=400, detail="Audio file is required for audio transcription.")
            try:
                contents = await request.audio_file.read()
                pipeline = model_data["pipeline"]
                try:
                    transcription = pipeline(contents, sampling_rate=16000)[0]["text"]  # Assuming 16kHz sampling rate
                    return JSONResponse({"transcription": transcription})
                except Exception as e:
                    raise HTTPException(status_code=500, detail=f"Error during audio transcription (pipeline): {str(e)}")
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Error during audio transcription (file read): {str(e)}")
        elif request.task_type == "summarization":
            if request.input_text is None:
                raise HTTPException(status_code=400, detail="Input text is required for summarization.")
            model = model_data["model"].to(device)
            tokenizer = model_data["tokenizer"]
            inputs = tokenizer(request.input_text, return_tensors="pt", truncation=True, max_length=512)  # added max_length for summarization
            with torch.no_grad():
                outputs = model.generate(**inputs)
            summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
            return JSONResponse({"summary": summary})

        else:
            raise HTTPException(status_code=500, detail=f"Unsupported task type")
    except Exception as e:
        logger.exception(f"Internal server error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")


@app.get("/", response_class=HTMLResponse)
async def root(request: Request):
    return TEMPLATES.TemplateResponse("index.html", {"request": request})

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

# Authentication Endpoints

@app.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
    user = authenticate_user(form_data.username, form_data.password)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(data={"sub": user["username"]}, expires_delta=access_token_expires)
    return {"access_token": access_token, "token_type": "bearer"}

def authenticate_user(username: str, password: str):
    user = get_user(username)
    if user and pwd_context.verify(password, user.hashed_password):
        return {"username": user.username}
    return None

def create_access_token(data: Dict[str, Any], expires_delta: timedelta = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

class Token(BaseModel):
    access_token: str
    token_type: str


@app.get("/users/me")
async def read_users_me(current_user: str = Depends(get_current_user)):
    return {"username": current_user}

async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = {"username": username, "token": token}
    except JWTError:
        raise credentials_exception
    user = get_user(username)
    if user is None:
        raise credentials_exception
    return username


@app.post("/register", response_model=User, status_code=status.HTTP_201_CREATED)
async def create_user(user: User):
    try:
        hashed_password = pwd_context.hash(user.password)
        new_user = {"username": user.username, "email": user.email, "hashed_password": hashed_password}
        inserted_user = insert_user(new_user)
        if inserted_user:
            return User(**inserted_user)
        else:
            raise HTTPException(status_code=500, detail="Failed to create user.")
    except Exception as e:
        logger.error(f"Error creating user: {e}")
        raise HTTPException(status_code=500, detail=f"Error creating user: {e}")


@app.put("/users/{username}", response_model=User, dependencies=[Depends(get_current_user)])
async def update_user_data(username: str, user: User):
    try:
        hashed_password = pwd_context.hash(user.password)
        updated_user_data = {"email": user.email, "hashed_password": hashed_password}
        updated_user = update_user(username, updated_user_data)
        if updated_user:
            return User(**updated_user)
        else:
            raise HTTPException(status_code=404, detail="User not found")

    except Exception as e:
        logger.error(f"Error updating user: {e}")
        raise HTTPException(status_code=500, detail="Error updating user.")



@app.delete("/users/{username}", dependencies=[Depends(get_current_user)])
async def delete_user_account(username: str):
    try:
        deleted_user = delete_user(username)
        if deleted_user:
            return JSONResponse({"message": "User deleted successfully."}, status_code=200)
        else:
            raise HTTPException(status_code=404, detail="User not found")
    except Exception as e:
        logger.error(f"Error deleting user: {e}")
        raise HTTPException(status_code=500, detail="Error deleting user.")


@app.get("/users", dependencies=[Depends(get_current_user)])
async def get_all_users_route():
    return get_all_users()



@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    return JSONResponse(
        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
        content=json.dumps({"detail": exc.errors(), "body": exc.body}),
    )


if __name__ == "__main__":

    create_db_and_table() # Initialize database on startup

    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)