main.py

import os
import torch
from imagebind import data
from imagebind.models import imagebind_model
from imagebind.models.imagebind_model import ModalityType
from pydub import AudioSegment
from fastapi import FastAPI, UploadFile, File, Form
from typing import List, Dict
import tempfile
from pydantic import BaseModel
import uvicorn
import numpy as np
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi import Depends, HTTPException, status

app = FastAPI()

# Add these lines after the app initialization
security = HTTPBearer()
API_TOKEN = os.getenv("API_TOKEN", "your-default-token-here")  # Set a default token or use environment variable

# Add this function for token verification
async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    if credentials.credentials != API_TOKEN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid authentication token",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return credentials.credentials

def convert_audio_to_wav(audio_path: str) -> str:
    """Convert MP3 to WAV if necessary."""
    if audio_path.lower().endswith('.mp3'):
        wav_path = audio_path.rsplit('.', 1)[0] + '.wav'
        if not os.path.exists(wav_path):
            audio = AudioSegment.from_mp3(audio_path)
            audio.export(wav_path, format='wav')
        return wav_path
    return audio_path

class EmbeddingManager:
    def __init__(self):
        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
        self.model = imagebind_model.imagebind_huge(pretrained=True)
        self.model.eval()
        self.model.to(self.device)
    
    def compute_embeddings(self, 
                          images: List[str] = None, 
                          audio_files: List[str] = None, 
                          texts: List[str] = None) -> dict:
        """Compute embeddings for provided modalities only."""
        with torch.no_grad():
            inputs = {}
            
            if texts:
                inputs[ModalityType.TEXT] = data.load_and_transform_text(texts, self.device)
            if images:
                inputs[ModalityType.VISION] = data.load_and_transform_vision_data(images, self.device)
            if audio_files:
                inputs[ModalityType.AUDIO] = data.load_and_transform_audio_data(audio_files, self.device)
            
            if not inputs:
                return {}
                
            embeddings = self.model(inputs)
            
            result = {}
            if ModalityType.VISION in inputs:
                result['vision'] = embeddings[ModalityType.VISION].cpu().numpy().tolist()
            if ModalityType.AUDIO in inputs:
                result['audio'] = embeddings[ModalityType.AUDIO].cpu().numpy().tolist()
            if ModalityType.TEXT in inputs:
                result['text'] = embeddings[ModalityType.TEXT].cpu().numpy().tolist()
                
            return result

    @staticmethod
    def compute_similarities(embeddings: Dict[str, List[List[float]]]) -> dict:
        """Compute similarities between available embeddings."""
        similarities = {}
        
        # Convert available embeddings to tensors
        tensors = {
            k: torch.tensor(v) for k, v in embeddings.items()
            if isinstance(v, (list, np.ndarray)) and len(v) > 0
        }
        
        # Compute cross-modal similarities
        modality_pairs = [
            ('vision', 'audio', 'vision_audio'),
            ('vision', 'text', 'vision_text'),
            ('audio', 'text', 'audio_text')
        ]
        
        for mod1, mod2, key in modality_pairs:
            if mod1 in tensors and mod2 in tensors:
                similarities[key] = torch.softmax(
                    tensors[mod1] @ tensors[mod2].T, 
                    dim=-1
                ).numpy().tolist()
        
        # Compute same-modality similarities
        for modality in ['vision', 'audio', 'text']:
            if modality in tensors:
                key = f'{modality}_{modality}'
                similarities[key] = torch.softmax(
                    tensors[modality] @ tensors[modality].T,
                    dim=-1
                ).numpy().tolist()
        
        return similarities

# Initialize the embedding manager
embedding_manager = EmbeddingManager()

class EmbeddingResponse(BaseModel):
    embeddings: dict
    file_names: dict

class SimilarityRequest(BaseModel):
    embeddings: Dict[str, List[List[float]]]
    threshold: float = 0.5
    top_k: int | None = None
    include_self_similarity: bool = False
    normalize_scores: bool = True

class SimilarityMatch(BaseModel):
    index_a: int
    index_b: int
    score: float
    modality_a: str
    modality_b: str
    item_a: str  # Original item identifier (filename or text)
    item_b: str  # Original item identifier (filename or text)

class SimilarityResponse(BaseModel):
    matches: List[SimilarityMatch]
    statistics: Dict[str, float]  # Contains avg_score, max_score, etc.
    modality_pairs: List[str]  # Lists which modality comparisons were performed

class ModalityPair:
    def __init__(self, mod1: str, mod2: str):
        self.mod1 = min(mod1, mod2)  # Ensure consistent ordering
        self.mod2 = max(mod1, mod2)
    
    def __str__(self):
        return f"{self.mod1}_to_{self.mod2}"

def compute_similarity_matrix(tensor1: torch.Tensor, tensor2: torch.Tensor, normalize: bool = True) -> torch.Tensor:
    """Compute cosine similarity between two sets of embeddings."""
    # Normalize embeddings if requested
    if normalize:
        tensor1 = torch.nn.functional.normalize(tensor1, dim=1)
        tensor2 = torch.nn.functional.normalize(tensor2, dim=1)
    
    # Compute similarity matrix
    similarity = torch.matmul(tensor1, tensor2.T)
    
    return similarity

def get_top_k_matches(similarity_matrix: torch.Tensor, top_k: int | None = None) -> List[tuple]:
    """Get top-k matches from a similarity matrix."""
    if top_k is None:
        top_k = similarity_matrix.numel()
    
    # Flatten and get top-k indices
    flat_sim = similarity_matrix.flatten()
    top_k = min(top_k, flat_sim.numel())
    values, indices = torch.topk(flat_sim, k=top_k)
    
    # Convert flat indices to 2D indices
    rows = indices // similarity_matrix.size(1)
    cols = indices % similarity_matrix.size(1)
    
    return [(r.item(), c.item(), v.item()) for r, c, v in zip(rows, cols, values)]

@app.post("/compute_embeddings", response_model=EmbeddingResponse)
async def generate_embeddings(
    credentials: HTTPAuthorizationCredentials = Depends(verify_token),
    texts: str | None = Form(None),
    images: List[UploadFile] | None = File(default=None),
    audio_files: List[UploadFile] | None = File(default=None)
):
    """Generate embeddings for any provided files and texts."""
    temp_files = []
    
    try:
        image_paths = []
        image_names = []
        audio_paths = []
        audio_names = []
        text_list = []
        
        # Process images if provided
        if images:
            for img in images:
                with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(img.filename)[1]) as tmp:
                    content = await img.read()
                    tmp.write(content)
                    image_paths.append(tmp.name)
                    image_names.append(img.filename)
                    temp_files.append(tmp.name)
        
        # Process audio files if provided
        if audio_files:
            for audio in audio_files:
                with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(audio.filename)[1]) as tmp:
                    content = await audio.read()
                    tmp.write(content)
                    audio_path = convert_audio_to_wav(tmp.name)
                    audio_paths.append(audio_path)
                    audio_names.append(audio.filename)
                    temp_files.append(tmp.name)
                    if audio_path != tmp.name:
                        temp_files.append(audio_path)
        
        # Process texts if provided
        if texts:
            text_list = [text.strip() for text in texts.split('\n') if text.strip()]
        
        # Compute embeddings only if we have any input
        if not any([image_paths, audio_paths, text_list]):
            return EmbeddingResponse(
                embeddings={},
                file_names={}
            )
        
        embeddings = embedding_manager.compute_embeddings(
            image_paths if image_paths else None,
            audio_paths if audio_paths else None,
            text_list if text_list else None
        )
        
        file_names = {}
        if image_names:
            file_names['images'] = image_names
        if audio_names:
            file_names['audio'] = audio_names
        if text_list:
            file_names['texts'] = text_list
        
        return EmbeddingResponse(
            embeddings=embeddings,
            file_names=file_names
        )
        
    finally:
        # Clean up temporary files
        for temp_file in temp_files:
            try:
                os.unlink(temp_file)
            except:
                pass

@app.post("/compute_similarities", response_model=SimilarityResponse)
async def compute_similarities(
    request: SimilarityRequest,
    file_names: Dict[str, List[str]],  # Maps modality to list of file/text names
    credentials: HTTPAuthorizationCredentials = Depends(verify_token)
):
    """
    Compute cross-modal similarities with advanced filtering and matching options.
    
    Parameters:
    - embeddings: Dict mapping modality to embedding tensors
    - threshold: Minimum similarity score to include in results
    - top_k: Maximum number of matches to return (per modality pair)
    - include_self_similarity: Whether to include same-item comparisons
    - normalize_scores: Whether to normalize embeddings before comparison
    - file_names: Dict mapping modality to list of original file/text names
    """
    
    matches = []
    statistics = {
        "avg_score": 0.0,
        "max_score": 0.0,
        "min_score": 1.0,
        "total_comparisons": 0
    }
    
    # Convert embeddings to tensors
    tensors = {
        k: torch.tensor(v) for k, v in request.embeddings.items()
        if isinstance(v, (list, np.ndarray)) and len(v) > 0
    }
    
    modality_pairs = []
    all_scores = []
    
    # Get all possible modality pairs
    modalities = list(tensors.keys())
    for i, mod1 in enumerate(modalities):
        for mod2 in modalities[i:]:  # Include self-comparisons if requested
            if mod1 == mod2 and not request.include_self_similarity:
                continue
                
            pair = ModalityPair(mod1, mod2)
            modality_pairs.append(str(pair))
            
            # Compute similarity matrix
            sim_matrix = compute_similarity_matrix(
                tensors[mod1], 
                tensors[mod2],
                normalize=request.normalize_scores
            )
            
            # Get top matches
            top_matches = get_top_k_matches(sim_matrix, request.top_k)
            
            # Filter by threshold and create match objects
            for idx_a, idx_b, score in top_matches:
                if score < request.threshold:
                    continue
                    
                # Skip self-matches if not requested
                if mod1 == mod2 and idx_a == idx_b and not request.include_self_similarity:
                    continue
                
                matches.append(SimilarityMatch(
                    index_a=idx_a,
                    index_b=idx_b,
                    score=float(score),
                    modality_a=mod1,
                    modality_b=mod2,
                    item_a=file_names[mod1][idx_a],
                    item_b=file_names[mod2][idx_b]
                ))
                all_scores.append(score)
    
    # Compute statistics
    if all_scores:
        statistics.update({
            "avg_score": float(np.mean(all_scores)),
            "max_score": float(np.max(all_scores)),
            "min_score": float(np.min(all_scores)),
            "total_comparisons": len(all_scores)
        })
    
    # Sort matches by score in descending order
    matches.sort(key=lambda x: x.score, reverse=True)
    
    return SimilarityResponse(
        matches=matches,
        statistics=statistics,
        modality_pairs=modality_pairs
    )

@app.get("/health")
async def health_check(
    credentials: HTTPAuthorizationCredentials = Depends(verify_token)
):
    """Basic healthcheck endpoint that returns the status of the service."""
    return {
        "status": "healthy",
        "model_device": embedding_manager.device
    }

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