app.py

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

from transformers import ViTImageProcessor, ViTForImageClassification 
from PIL import Image
import io
import uvicorn
import os
import logging

#setup for basic logging 
logging.basicConfig(level=logging.INFO)

app = FastAPI()

# Global variables for models
classifier = None # Text Emotion Classifier
llm_pipeline = None
# **NEW GLOBAL VARIABLE FOR FER MODEL**
fer_model = None
fer_processor = None
fer_id2label = None

# --- Model Loading and Initialization ---
@app.on_event("startup")
async def load_model():
    global classifier, llm_pipeline, fer_model, fer_processor, fer_id2label
    
    logging.info("Starting model loading...")
    
    # 1. Load the Emotion Classifier (Text Classification)
    try:
        classifier = pipeline("text-classification", model="MilaNLProc/xlm-emo-t")
        logging.info("Text emotion classifier loaded successfully.")
    except Exception as e:
        logging.error(f"Failed to load text classifier: {e}")

    # 2. Load the Lightweight LLM (TinyLlama-1.1B-Chat-v1.0)
    try:
        model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" 
        model = AutoModelForCausalLM.from_pretrained(model_id)
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        
        llm_pipeline = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            max_new_tokens=80,
            temperature=0.6,
            do_sample=True,
            return_full_text=False
        )
        logging.info("LLM transformers pipeline initialized successfully.")
    except Exception as e:
        logging.error(f"Failed to initialize LLM pipeline: {e}")

    # **3. Load the Facial Emotion Recognition (FER) Model**
    try:
        fer_model_id = "dima806/facial_emotions_image_detection"
        fer_processor = ViTImageProcessor.from_pretrained(fer_model_id)
        fer_model = ViTForImageClassification.from_pretrained(fer_model_id)
        fer_id2label = fer_model.config.id2label
        logging.info("FER model loaded successfully.")
    except Exception as e:
        logging.error(f"Failed to load FER model: {e}")

# --- API Data Models ---
class TextRequest(BaseModel):
    text: str

class FeedbackRequest(BaseModel):
    text: str
    emotion: str

# --- API Endpoint: Text Emotion Classification ---
@app.post("/api/classify")
async def classify_emotion_api(request: TextRequest):
    if not request.text or not classifier:
        return {"emotion": "error", "score": 0.0}

    try:
        result = classifier(request.text)
        top = result[0]
        return {
            "emotion": top["label"],
            "score": round(top["score"] * 100, 2)
        }
    except Exception as e:
        logging.error(f"Error during classification: {e}")
        return {"emotion": "api_error", "score": 0.0}

# --- API Endpoint: Facial Emotion Recognition ---
@app.post("/api/classify_face")
async def classify_face_api(file: UploadFile = File(...)):
    if not fer_model or not fer_processor:
        raise HTTPException(status_code=500, detail="FER model not initialized.")

    try:
        # Read the image file and convert it to a PIL Image
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        
        # Preprocess and classify
        inputs = fer_processor(images=image, return_tensors="pt")
        outputs = fer_model(**inputs)
        logits = outputs.logits
        
        # Get the predicted class and confidence
        predicted_class_idx = logits.argmax(-1).item()
        emotion = fer_id2label[predicted_class_idx]
        
        # Calculate score (using softmax)
        import torch
        confidence_tensor = torch.nn.functional.softmax(logits, dim=-1)[0]
        score = confidence_tensor[predicted_class_idx].item() * 100

        return {
            "emotion": emotion,
            "score": round(score, 2)
        }

    except Exception as e:
        logging.error(f"Error during facial classification: {e}")
        raise HTTPException(status_code=500, detail="Error processing image for FER.")


# --- API Endpoint: LLM Feedback Generation ---
@app.post("/api/llm_feedback")
async def llm_feedback_api(req: FeedbackRequest):
    if not llm_pipeline:
        return {"feedback": "Error: LLM pipeline not initialized."}
        
    prompt = f"""
<|system|>
Give a direct, natural, supportive reply in one short sentence. 
Do not mention any instructions, do not refer to the user as "user", 
and do not describe what they said. Just respond directly.
<|user|>
{req.text}
Emotion: {req.emotion}
<|assistant|>
"""
    
    try:
        result = llm_pipeline(prompt)
        response = result[0]['generated_text'].strip()
        return {"feedback": response}
        
    except Exception as e:
        logging.error(f"LLM text generation failed: {e}")
        return {"feedback": f"I detected {req.emotion}. Take care."}


# Serve frontend 
app.mount("/", StaticFiles(directory=".", html=True), name="static")

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