Update app.py

app.py

@@ -1,1303 +1,474 @@
-import streamlit as st
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchaudio
-import torchaudio.transforms as T
-import torchvision.transforms as transforms
-import numpy as np
-import cv2
-from pathlib import Path
-from PIL import Image
-import plotly.graph_objects as go
-import plotly.express as px
-import io
-import timm
-import warnings
-import tempfile
-
-warnings.filterwarnings('ignore')
-
-# ======================= PAGE CONFIG ============================
-st.set_page_config(
-    page_title="DFDA - Deepfake Detection & Attribution",
-    page_icon="https://play-lh.googleusercontent.com/wJEqeRZGZds8tVZKQq-azZ1sp7IjYgF92zXVKYD1vy1LrV6xz4KVriABCjdjS8jQyA",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
-
-# ======================= CUSTOM CSS ============================
-st.markdown("""
-<style>
-    /* Import Google Fonts */
-    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700;800&display=swap');
-    
-    /* Global Styles */
-    * {
-        font-family: 'Inter', sans-serif;
-    }
-    
-    /* Hide Streamlit Branding */
-    #MainMenu {visibility: hidden;}
-    footer {visibility: hidden;}
-    header {visibility: hidden;}
-    
-    /* Main Container */
-    .main {
-        background: linear-gradient(135deg, #0f0f1e 0%, #1a1a2e 100%);
-    }
-    
-    /* Sidebar Styling */
-    [data-testid="stSidebar"] {
-        background: linear-gradient(180deg, #16213e 0%, #0f3460 100%);
-        border-right: 2px solid rgba(94, 129, 244, 0.3);
-    }
-    
-    [data-testid="stSidebar"] [data-testid="stMarkdownContainer"] {
-        color: #e4e4e7;
-    }
-    
-    /* Header Gradient */
-    .gradient-header {
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 50%, #f093fb 100%);
-        padding: 3rem 2rem;
-        border-radius: 10px;
-        text-align: center;
-        margin-bottom: 2rem;
-        box-shadow: 0 10px 60px rgba(102, 126, 234, 0.4);
-        animation: fadeInDown 0.8s ease-out;
-    }
-    
-    .gradient-header h1 {
-        color: white;
-        font-size: 2.5rem;
-        font-weight: 800;
-        margin: 0;
-        text-shadow: 2px 4px 8px rgba(0,0,0,0.3);
-        letter-spacing: -1px;
-    }
-    
-    .gradient-header p {
-        color: rgba(255,255,255,0.95);
-        font-size: 1.2rem;
-        margin-top: 0.8rem;
-        font-weight: 500;
-    }
-    
-    /* Feature Cards */
-    .feature-card {
-        background: linear-gradient(135deg, rgba(102, 126, 234, 0.15) 0%, rgba(118, 75, 162, 0.15) 100%);
-        border: 2px solid rgba(102, 126, 234, 0.3);
-        border-radius: 16px;
-        padding: 2rem;
-        margin: 1.5rem 0;
-        transition: all 0.4s cubic-bezier(0.4, 0, 0.2, 1);
-        backdrop-filter: blur(10px);
-        position: relative;
-        overflow: hidden;
-    }
-    
-    .feature-card::before {
-        content: '';
-        position: absolute;
-        top: 0;
-        left: -100%;
-        width: 100%;
-        height: 100%;
-        background: linear-gradient(90deg, transparent, rgba(255,255,255,0.1), transparent);
-        transition: left 0.5s;
-    }
-    
-    .feature-card:hover::before {
-        left: 100%;
-    }
-    
-    .feature-card:hover {
-        transform: translateY(-8px) scale(1.02);
-        box-shadow: 0 20px 40px rgba(102, 126, 234, 0.3);
-        border-color: rgba(102, 126, 234, 0.6);
-    }
-    
-    .feature-title {
-        color: #a5b4fc;
-        font-size: 1.4rem;
-        font-weight: 700;
-        margin-bottom: 1rem;
-        letter-spacing: -0.5px;
-    }
-    
-    .feature-desc {
-        color: #e4e4e7;
-        font-size: 1.05rem;
-        line-height: 1.8;
-    }
-    
-    /* Result Cards */
-    .result-success {
-        background: linear-gradient(135deg, rgba(16, 185, 129, 0.2) 0%, rgba(5, 150, 105, 0.2) 100%);
-        border-left: 6px solid #10b981;
-        border-radius: 16px;
-        padding: 2rem;
-        margin: 2rem 0;
-        box-shadow: 0 10px 30px rgba(16, 185, 129, 0.2);
-        animation: slideInLeft 0.6s ease-out;
-    }
-    
-    .result-warning {
-        background: linear-gradient(135deg, rgba(239, 68, 68, 0.2) 0%, rgba(220, 38, 38, 0.2) 100%);
-        border-left: 6px solid #ef4444;
-        border-radius: 16px;
-        padding: 2rem;
-        margin: 2rem 0;
-        box-shadow: 0 10px 30px rgba(239, 68, 68, 0.2);
-        animation: slideInLeft 0.6s ease-out;
-    }
-    
-    .result-title {
-        font-size: 2rem;
-        font-weight: 800;
-        margin: 0 0 1rem 0;
-        letter-spacing: -0.5px;
-    }
-    
-    .result-subtitle {
-        font-size: 1.3rem;
-        font-weight: 600;
-        margin: 0.8rem 0;
-        opacity: 0.95;
-    }
-    
-    /* Metric Cards */
-    .metric-card {
-        background: linear-gradient(135deg, rgba(102, 126, 234, 0.1) 0%, rgba(118, 75, 162, 0.1) 100%);
-        border: 2px solid rgba(102, 126, 234, 0.3);
-        border-radius: 16px;
-        padding: 2rem;
-        text-align: center;
-        margin: 1rem 0.5rem;
-        transition: all 0.3s ease;
-        backdrop-filter: blur(10px);
-    }
-    
-    .metric-card:hover {
-        transform: translateY(-5px);
-        box-shadow: 0 15px 35px rgba(102, 126, 234, 0.3);
-    }
-    
-    .metric-value {
-        font-size: 3rem;
-        font-weight: 800;
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-        -webkit-background-clip: text;
-        -webkit-text-fill-color: transparent;
-        background-clip: text;
-        margin-bottom: 0.5rem;
-    }
-    
-    .metric-label {
-        font-size: 1rem;
-        color: #a5b4fc;
-        font-weight: 600;
-        text-transform: uppercase;
-        letter-spacing: 1.5px;
-    }
-    
-    /* Info Boxes */
-    .info-box {
-        background: linear-gradient(135deg, rgba(59, 130, 246, 0.15) 0%, rgba(37, 99, 235, 0.15) 100%);
-        border-left: 5px solid #3b82f6;
-        border-radius: 12px;
-        padding: 1.5rem;
-        margin: 1.5rem 0;
-        backdrop-filter: blur(10px);
-    }
-    
-    .info-box h3 {
-        color: #93c5fd;
-        font-size: 1.4rem;
-        font-weight: 700;
-        margin: 0 0 0.8rem 0;
-    }
-    
-    .info-box p {
-        color: #e4e4e7;
-        font-size: 1.05rem;
-        line-height: 1.7;
-        margin: 0;
-    }
-    
-    /* Buttons */
-    .stButton > button {
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-        color: white;
-        border: none;
-        border-radius: 12px;
-        padding: 1rem 3rem;
-        font-weight: 700;
-        font-size: 1.1rem;
-        transition: all 0.3s ease;
-        box-shadow: 0 8px 20px rgba(102, 126, 234, 0.4);
-        text-transform: uppercase;
-        letter-spacing: 1px;
-    }
-    
-    .stButton > button:hover {
-        transform: translateY(-3px);
-        box-shadow: 0 12px 30px rgba(102, 126, 234, 0.5);
-        background: linear-gradient(135deg, #764ba2 0%, #667eea 100%);
-    }
-    
-    /* File Uploader */
-    [data-testid="stFileUploader"] {
-        border: 3px dashed rgba(102, 126, 234, 0.5);
-        border-radius: 16px;
-        padding: 2rem;
-        background: rgba(102, 126, 234, 0.05);
-        transition: all 0.3s ease;
-    }
-    
-    [data-testid="stFileUploader"]:hover {
-        border-color: rgba(102, 126, 234, 0.8);
-        background: rgba(102, 126, 234, 0.1);
-    }
-    
-    /* Progress Bar */
-    .stProgress > div > div > div > div {
-        background: linear-gradient(90deg, #667eea 0%, #764ba2 50%, #f093fb 100%);
-    }
-    
-    /* Radio Buttons */
-    [data-testid="stRadio"] > div {
-        background: rgba(102, 126, 234, 0.1);
-        border-radius: 12px;
-        padding: 1rem;
-    }
-    
-    /* Divider */
-    hr {
-        border: none;
-        height: 2px;
-        background: linear-gradient(90deg, transparent, rgba(102, 126, 234, 0.5), transparent);
-        margin: 2rem 0;
-    }
-    
-    /* Animations */
-    @keyframes fadeInDown {
-        from {
-            opacity: 0;
-            transform: translateY(-30px);
-        }
-        to {
-            opacity: 1;
-            transform: translateY(0);
-        }
-    }
-    
-    @keyframes slideInLeft {
-        from {
-            opacity: 0;
-            transform: translateX(-30px);
-        }
-        to {
-            opacity: 1;
-            transform: translateX(0);
-        }
-    }
-    
-    @keyframes pulse {
-        0%, 100% {
-            opacity: 1;
-        }
-        50% {
-            opacity: 0.7;
-        }
-    }
-    
-    /* Expander */
-    [data-testid="stExpander"] {
-        background: rgba(102, 126, 234, 0.1);
-        border-radius: 12px;
-        border: 1px solid rgba(102, 126, 234, 0.3);
-    }
-    
-    /* Slider */
-    .stSlider > div > div > div {
-        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
-    }
-</style>
-""", unsafe_allow_html=True)
-
-# ======================= CONFIG CLASSES ============================
-class AudioConfig:
-    MODEL_PATH = Path("Audio.pth")
-    CLASS_NAMES = ["ElevenLabs", "Original", "Tacotron", "Text To Speech", "Voice Conversion"]
-    SAMPLE_RATE = 16000
-    N_MELS = 128
-    N_FFT = 2048
-    HOP_LENGTH = 512
-    MAX_AUDIO_LENGTH = 5
-    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-class ImageConfig:
-    MODEL_PATH = Path("Image.pt")
-    CLASS_NAMES = ["DALL-E", "DeepFaceLab", "Face2Face", "FaceShifter", "FaceSwap", "Midjourney", 
-                   "NeuralTextures", "Real", "Stable Diffusion", "StyleGAN"]
-    IMG_SIZE = 299
-    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-class VideoConfig:
-    MODEL_PATH = Path("Video.pt")
-    CLASS_NAMES = ["Deepfakes", "Face2Face", "FaceShifter", "FaceSwap", "NeuralTextures", "Original"]
-    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    FRAME_SKIP = 10
-
-# ======================= AUDIO MODEL ============================
-class AudioDeepfakeCNN(nn.Module):
-    def __init__(self, num_classes=5):
-        super().__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(1, 64, 3, padding=1),
-            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
-            nn.Conv2d(64, 64, 3, padding=1),
-            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
-            nn.MaxPool2d(2,2),
-            nn.Dropout2d(0.25)
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(64, 128, 3, padding=1),
-            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
-            nn.Conv2d(128, 128, 3, padding=1),
-            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
-            nn.MaxPool2d(2,2),
-            nn.Dropout2d(0.25)
-        )
-        self.conv3 = nn.Sequential(
-            nn.Conv2d(128, 256, 3, padding=1),
-            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
-            nn.Conv2d(256, 256, 3, padding=1),
-            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
-            nn.MaxPool2d(2,2),
-            nn.Dropout2d(0.3)
-        )
-        self.conv4 = nn.Sequential(
-            nn.Conv2d(256, 512, 3, padding=1),
-            nn.BatchNorm2d(512), nn.ReLU(inplace=True),
-            nn.Conv2d(512, 512, 3, padding=1),
-            nn.BatchNorm2d(512), nn.ReLU(inplace=True),
-            nn.MaxPool2d(2,2),
-            nn.Dropout2d(0.3)
-        )
-        self.gap = nn.AdaptiveAvgPool2d((1,1))
-        self.fc = nn.Sequential(
-            nn.Linear(512, 256), nn.ReLU(inplace=True), nn.Dropout(0.5),
-            nn.Linear(256, 128), nn.ReLU(inplace=True), nn.Dropout(0.4),
-            nn.Linear(128, num_classes)
-        )
-    def forward(self, x):
-        x = self.conv1(x); x = self.conv2(x); x = self.conv3(x); x = self.conv4(x)
-        x = self.gap(x); x = x.view(x.size(0), -1); x = self.fc(x); return x
-
-# ======================= IMAGE MODEL ============================
-class AttentionBlock(nn.Module):
-    def __init__(self, in_features):
-        super(AttentionBlock, self).__init__()
-        self.attention = nn.Sequential(
-            nn.Linear(in_features, in_features // 8),
-            nn.ReLU(),
-            nn.Linear(in_features // 8, in_features),
-            nn.Sigmoid()
-        )
-    def forward(self, x):
-        att_weights = self.attention(x)
-        return x * att_weights
-
-class EnhancedXceptionNet(nn.Module):
-    def __init__(self, num_classes, dropout_rate=0.3, use_attention=True):
-        super(EnhancedXceptionNet, self).__init__()
-        self.base_model = timm.create_model('xception', pretrained=False)
-        in_features = self.base_model.get_classifier().in_features
-        self.base_model.reset_classifier(0)
-        self.use_attention = use_attention
-        if use_attention:
-            self.attention = AttentionBlock(in_features)
-        self.classifier = nn.Sequential(
-            nn.Dropout(dropout_rate),
-            nn.Linear(in_features, 512),
-            nn.BatchNorm1d(512),
-            nn.ReLU(),
-            nn.Dropout(dropout_rate / 2),
-            nn.Linear(512, num_classes)
-        )
-    def forward(self, x):
-        features = self.base_model(x)
-        if self.use_attention:
-            features = self.attention(features)
-        output = self.classifier(features)
-        return output
-
-# ======================= MODEL LOADERS ============================
-@st.cache_resource
-def load_audio_model():
-    try:
-        model = AudioDeepfakeCNN(num_classes=len(AudioConfig.CLASS_NAMES))
-        checkpoint = torch.load(AudioConfig.MODEL_PATH, map_location=AudioConfig.DEVICE)
-        model.load_state_dict(checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint)
-        model.to(AudioConfig.DEVICE).eval()
-        return model
-    except Exception as e:
-        st.error(f"Error loading audio model: {e}")
-        return None
-
-@st.cache_resource
-def load_image_model():
-    try:
-        model = EnhancedXceptionNet(num_classes=len(ImageConfig.CLASS_NAMES), dropout_rate=0.3, use_attention=True)
-        checkpoint = torch.load(ImageConfig.MODEL_PATH, map_location=ImageConfig.DEVICE)
-        state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
-        model_dict = model.state_dict()
-        filtered_dict = {k: v for k, v in state_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
-        model_dict.update(filtered_dict)
-        model.load_state_dict(model_dict, strict=False)
-        model.to(ImageConfig.DEVICE).eval()
-        return model
-    except Exception as e:
-        st.error(f"Error loading image model: {e}")
-        return None
-
-@st.cache_resource
-def load_video_model():
-    try:
-        num_classes = len(VideoConfig.CLASS_NAMES)
-        model = timm.create_model("legacy_xception", pretrained=False, num_classes=num_classes)
-        checkpoint = torch.load(VideoConfig.MODEL_PATH, map_location=VideoConfig.DEVICE)
-        state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
-        model_dict = model.state_dict()
-        filtered_dict = {k: v for k, v in state_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
-        model_dict.update(filtered_dict)
-        model.load_state_dict(model_dict, strict=False)
-        model.to(VideoConfig.DEVICE).eval()
-        return model
-    except Exception as e:
-        st.error(f"Error loading video model: {e}")
-        return None
-
-# ======================= PREPROCESSING ============================
-def preprocess_audio(audio_path, sample_rate, n_mels, n_fft, hop_length, max_length):
-    waveform, sr = torchaudio.load(str(audio_path))
-    if waveform.shape[0] > 1:
-        waveform = torch.mean(waveform, dim=0, keepdim=True)
-    if sr != sample_rate:
-        resampler = T.Resample(sr, sample_rate)
-        waveform = resampler(waveform)
-    max_len = sample_rate * max_length
-    if waveform.shape[1] < max_len:
-        waveform = torch.nn.functional.pad(waveform, (0, max_len - waveform.shape[1]))
-    else:
-        waveform = waveform[:, :max_len]
-    mel_spec = T.MelSpectrogram(sample_rate, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels, power=2.0)(waveform)
-    mel_spec_db = T.AmplitudeToDB(top_db=80)(mel_spec)
-    mel_spec_db = (mel_spec_db - mel_spec_db.mean()) / (mel_spec_db.std() + 1e-8)
-    return mel_spec_db, mel_spec_db.squeeze().numpy()
-
-def preprocess_image(image_path, size):
-    trans = transforms.Compose([
-        transforms.Resize((size, size)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
-    ])
-    img = Image.open(image_path).convert('RGB')
-    return trans(img).unsqueeze(0)
-
-def extract_frames(video_path, frame_skip=10):
-    frames = []
-    cap = cv2.VideoCapture(str(video_path))
-    count = 0
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret: break
-        if count % frame_skip == 0:
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            frames.append(frame)
-        count += 1
-    cap.release()
-    return frames
-
-def preprocess_frame(frame, size):
-    trans = transforms.Compose([
-        transforms.ToPILImage(),
-        transforms.Resize((size, size)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
-    ])
-    return trans(frame).unsqueeze(0)
-
-# ======================= INFERENCE ============================
-def predict_audio(model, audio_path, config):
-    device = config.DEVICE
-    try:
-        x, spec_viz = preprocess_audio(audio_path, config.SAMPLE_RATE, config.N_MELS, config.N_FFT, config.HOP_LENGTH, config.MAX_AUDIO_LENGTH)
-        x = x.unsqueeze(0).to(device)
-        with torch.no_grad():
-            out = F.softmax(model(x), dim=1)[0]
-        pred_idx = torch.argmax(out).item()
-        return {
-            "predicted_class": config.CLASS_NAMES[pred_idx],
-            "predicted_index": pred_idx,
-            "confidence": out[pred_idx].item(),
-            "probabilities": {c: out[i].item() for i, c in enumerate(config.CLASS_NAMES)}
-        }, spec_viz
-    except Exception as e:
-        return {"error": str(e)}, None
-
-def predict_image(model, image_path, config):
-    device = config.DEVICE
-    try:
-        x = preprocess_image(image_path, config.IMG_SIZE).to(device)
-        with torch.no_grad():
-            out = F.softmax(model(x), dim=1)[0]
-        pred_idx = torch.argmax(out).item()
-        return {
-            "predicted_class": config.CLASS_NAMES[pred_idx],
-            "predicted_index": pred_idx,
-            "confidence": out[pred_idx].item(),
-            "probabilities": {c: out[i].item() for i, c in enumerate(config.CLASS_NAMES)}
-        }
-    except Exception as e:
-        return {"error": str(e)}
-
-def predict_video(model, video_path, config):
-    device = config.DEVICE
-    try:
-        frames = extract_frames(video_path, config.FRAME_SKIP)
-        if len(frames) == 0:
-            return {"error":"No frames extracted"}
-        probs = []
-        progress_bar = st.progress(0)
-        for i, frame in enumerate(frames):
-            x = preprocess_frame(frame, 300).to(device)
-            with torch.no_grad():
-                p = F.softmax(model(x),dim=1)[0].cpu().numpy()
-            probs.append(p)
-            progress_bar.progress((i + 1) / len(frames))
-        probs = np.stack(probs)
-        avg_prob = probs.mean(axis=0)
-        pred_idx = np.argmax(avg_prob)
-        return {
-            "predicted_class": config.CLASS_NAMES[pred_idx],
-            "predicted_index": pred_idx,
-            "confidence": float(avg_prob[pred_idx]),
-            "probabilities": {c: float(avg_prob[i]) for i, c in enumerate(config.CLASS_NAMES)},
-            "frame_count": len(frames)
-        }
-    except Exception as e:
-        return {"error": str(e)}
-
-# ======================= VISUALIZATION ============================
-def plot_probabilities(probs, class_names):
-    fig = go.Figure(data=[
-        go.Bar(
-            x=list(probs.values()),
-            y=list(probs.keys()),
-            orientation='h',
-            marker=dict(
-                color=list(probs.values()),
-                colorscale=[[0, '#667eea'], [0.5, '#764ba2'], [1, '#f093fb']],
-                line=dict(color='rgba(102, 126, 234, 0.8)', width=2)
-            ),
-            text=[f'{v*100:.2f}%' for v in probs.values()],
-            textposition='outside',
-            textfont=dict(size=14, color='#e4e4e7', family='Inter')
-        )
-    ])
-    
-    fig.update_layout(
-        title=dict(text="Class Probability Distribution", font=dict(size=22, color='#a5b4fc', family='Inter')),
-        xaxis_title="Probability",
-        yaxis_title="Class",
-        height=400,
-        template="plotly_dark",
-        showlegend=False,
-        plot_bgcolor='rgba(15,15,30,0.5)',
-        paper_bgcolor='rgba(0,0,0,0)',
-        font=dict(color='#e4e4e7', family='Inter'),
-        xaxis=dict(gridcolor='rgba(102, 126, 234, 0.2)'),
-        yaxis=dict(gridcolor='rgba(102, 126, 234, 0.2)')
-    )
-    
-    return fig
-
-def plot_spectrogram(spec_data):
-    if spec_data is None:
-        return None
-    
-    fig = px.imshow(
-        spec_data,
-        labels=dict(x="Time", y="Frequency", color="Amplitude (dB)"),
-        aspect='auto',
-        color_continuous_scale=[[0, '#0f0f1e'], [0.25, '#667eea'], [0.5, '#764ba2'], [0.75, '#f093fb'], [1, '#fbbf24']]
-    )
-    
-    fig.update_layout(
-        title=dict(text="Mel Spectrogram Visualization", font=dict(size=22, color='#a5b4fc', family='Inter')),
-        height=400,
-        template="plotly_dark",
-        plot_bgcolor='rgba(15,15,30,0.5)',
-        paper_bgcolor='rgba(0,0,0,0)',
-        font=dict(color='#e4e4e7', family='Inter')
-    )
-    
-    return fig
-
-def plot_confusion_matrix(probs, predicted_class, class_names):
-    n_classes = len(class_names)
-    cm = np.zeros((n_classes, n_classes))
-    pred_idx = class_names.index(predicted_class)
-    for i, (cls, prob) in enumerate(probs.items()):
-        cm[pred_idx, i] = prob
-    
-    fig = px.imshow(
-        cm,
-        labels=dict(x="Predicted Class", y="True Class", color="Probability"),
-        x=class_names,
-        y=class_names,
-        color_continuous_scale=[[0, '#0f0f1e'], [0.5, '#667eea'], [1, '#f093fb']],
-        aspect='auto',
-        text_auto='.2f'
-    )
-    
-    fig.update_layout(
-        title=dict(text="Prediction Confidence Matrix", font=dict(size=22, color='#a5b4fc', family='Inter')),
-        height=500,
-        template="plotly_dark",
-        plot_bgcolor='rgba(15,15,30,0.5)',
-        paper_bgcolor='rgba(0,0,0,0)',
-        font=dict(color='#e4e4e7', family='Inter')
-    )
-    
-    return fig
-
-# ======================= PAGES ============================
-def show_home_page():
-    st.markdown("""
-    <div class="gradient-header">
-        <h1>Deepfake Detection & Attribution System</h1>
-        <p>Advanced AI-Powered Media Forensics Platform</p>
-    </div>
-    """, unsafe_allow_html=True)
-    
-    st.markdown("## Comprehensive Multi-Modal Deepfake Analysis System")
-    st.markdown("This cutting-edge research platform leverages state-of-the-art deep learning architectures to detect and attribute synthetic media across audio, image, and video modalities with unparalleled accuracy.")
-    
-    col1, col2, col3 = st.columns(3)
-    
-    with col1:
-        st.markdown("""
-        <div class="feature-card">
-            <div class="feature-title">Audio Deepfake Detection</div>
-            <div class="feature-desc">
-                Utilizes advanced mel-spectrogram analysis combined with custom CNN architecture to identify synthetic audio generation techniques. Our model detects voice cloning, text-to-speech systems, and voice conversion with exceptional precision.
-                <br><br>
-                <strong>Supported Formats:</strong> WAV, MP3, OGG, FLAC, M4A<br>
-                <strong>Detection Classes:</strong> ElevenLabs, Tacotron, Voice Conversion, Original Audio, Text-to-Speech<br>
-                <strong>Accuracy:</strong> 94.7% on validation set
-            </div>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    with col2:
-        st.markdown("""
-        <div class="feature-card">
-            <div class="feature-title">Image Synthesis Detection</div>
-            <div class="feature-desc">
-                Employs Enhanced XceptionNet with attention mechanisms to detect AI-generated and manipulated images. Identifies various deepfake techniques and distinguishes between multiple generative models.
-                <br><br>
-                <strong>Supported Formats:</strong> JPG, PNG, BMP, WEBP<br>
-                <strong>Detection Classes:</strong> DALL-E, Midjourney, StyleGAN, FaceSwap, DeepFaceLab, Real Images<br>
-                <strong>Accuracy:</strong> 96.2% on benchmark datasets
-            </div>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    with col3:
-        st.markdown("""
-        <div class="feature-card">
-            <div class="feature-title">Video Manipulation Analysis</div>
-            <div class="feature-desc">
-                Performs temporal consistency analysis across video frames using Xception-based architecture. Detects face manipulation, deepfake videos, and various facial reenactment techniques through frame-by-frame analysis.
-                <br><br>
-                <strong>Supported Formats:</strong> MP4, AVI, MOV, MKV, WEBM<br>
-                <strong>Detection Classes:</strong> Deepfakes, Face2Face, FaceSwap, FaceShifter, Original Video<br>
-                <strong>Accuracy:</strong> 93.8% with temporal aggregation
-            </div>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    st.markdown("---")
-    st.markdown("## System Workflow")
-    
-    col1, col2, col3 = st.columns(3)
-    with col1:
-        st.markdown("""
-        <div class="info-box">
-            <h3>STEP 1: Upload Media</h3>
-            <p>Select your analysis mode from the sidebar navigation menu. Upload your media file using the intuitive drag-and-drop interface or browse your local filesystem. The system supports multiple file formats across all modalities.</p>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    with col2:
-        st.markdown("""
-        <div class="info-box">
-            <h3>STEP 2: Deep Analysis</h3>
-            <p>Our advanced neural networks process your media through multiple layers of feature extraction and classification. The models analyze spectral features, pixel patterns, and temporal inconsistencies to detect manipulation.</p>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    with col3:
-        st.markdown("""
-        <div class="info-box">
-            <h3>STEP 3: Comprehensive Results</h3>
-            <p>Receive detailed analysis including predicted class, confidence scores, probability distributions, and visual representations. Export results for further investigation or documentation purposes.</p>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    st.markdown("---")
-    st.markdown("## Technical Specifications")
-    
-    col1, col2, col3, col4 = st.columns(4)
-    
-    with col1:
-        device = "CUDA GPU" if torch.cuda.is_available() else "CPU"
-        device_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
-        st.markdown(f"""
-        <div class="metric-card">
-            <div class="metric-value">{device}</div>
-            <div class="metric-label">Computation Device</div>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    with col2:
-        st.markdown("""
-        <div class="metric-card">
-            <div class="metric-value">3</div>
-            <div class="metric-label">Modalities</div>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    with col3:
-        total_classes = len(AudioConfig.CLASS_NAMES) + len(ImageConfig.CLASS_NAMES) + len(VideoConfig.CLASS_NAMES)
-        st.markdown(f"""
-        <div class="metric-card">
-            <div class="metric-value">{total_classes}</div>
-            <div class="metric-label">Detection Classes</div>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    with col4:
-        st.markdown("""
-        <div class="metric-card">
-            <div class="metric-value">92.5%</div>
-            <div class="metric-label">Avg Accuracy</div>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    st.markdown("---")
-    
-    with st.expander("Model Architecture Details"):
-        st.markdown("""
-        **Audio Detection Model:**
-        - Architecture: Custom 4-layer CNN with batch normalization and dropout
-        - Input: Mel spectrograms (128 mel bins, 5-second audio clips)
-        - Parameters: ~2.1M trainable parameters
-        - Training: Adam optimizer, learning rate 0.0001
-        
-        **Image Detection Model:**
-        - Architecture: Enhanced XceptionNet with attention mechanism
-        - Input: RGB images resized to 299x299 pixels
-        - Parameters: ~22.8M trainable parameters
-        - Training: SGD optimizer with momentum 0.9
-        
-        **Video Detection Model:**
-        - Architecture: Legacy Xception with temporal aggregation
-        - Input: Frame sequences extracted at configurable intervals
-        - Parameters: ~22.9M trainable parameters
-        - Training: Frame-level predictions with temporal averaging
-        """)
-    
-    with st.expander("Research & Citations"):
-        st.markdown("""
-        This system is built upon cutting-edge research in deepfake detection and media forensics:
-        
-        - XceptionNet architecture for image manipulation detection
-        - Mel-spectrogram analysis for synthetic audio detection
-        - Temporal consistency analysis for video deepfake detection
-        - Attention mechanisms for improved feature selection
-        
-        **Recommended for:**
-        - Academic research in digital forensics
-        - Media verification and fact-checking
-        - Content authenticity assessment
-        - Deepfake dataset curation and analysis
-        """)
-
-def show_audio_page():
-    st.markdown("""
-    <div class="gradient-header">
-        <h1>AUDIO DEEPFAKE DETECTION</h1>
-        <p>Synthetic Voice Generation & Manipulation Analysis</p>
-    </div>
-    """, unsafe_allow_html=True)
-    
-    model = load_audio_model()
-    if model is None:
-        st.error("Audio model not loaded. Please ensure Audio.pth is in the correct directory.")
-        return
-    
-    st.markdown("### Upload Audio File for Analysis")
-    uploaded_file = st.file_uploader(
-        "Drag and drop your audio file here or click to browse",
-        type=['wav', 'mp3', 'ogg', 'flac', 'm4a'],
-        help="Supported formats: WAV, MP3, OGG, FLAC, M4A (Max 5 seconds will be analyzed)"
-    )
-    
-    if uploaded_file is not None:
-        col1, col2 = st.columns([1, 1])
-        
-        with col1:
-            st.markdown("#### Audio Preview")
-            st.audio(uploaded_file, format='audio/wav')
-            st.markdown(f"""
-            **File Information:**
-            - Filename: `{uploaded_file.name}`
-            - Size: `{uploaded_file.size / 1024:.2f} KB`
-            - Type: `{uploaded_file.type}`
-            """)
-        
-        with col2:
-            st.markdown("#### Analysis Configuration")
-            st.markdown(f"""
-            **Model Settings:**
-            - Sample Rate: {AudioConfig.SAMPLE_RATE} Hz
-            - Mel Bins: {AudioConfig.N_MELS}
-            - Analysis Duration: {AudioConfig.MAX_AUDIO_LENGTH} seconds
-            - Device: {AudioConfig.DEVICE}
-            """)
-            
-            if st.button("ANALYZE AUDIO", type="primary", use_container_width=True):
-                with st.spinner("Processing audio file and extracting features..."):
-                    with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp_file:
-                        tmp_file.write(uploaded_file.getvalue())
-                        tmp_path = tmp_file.name
-                    
-                    result, spec_viz = predict_audio(model, Path(tmp_path), AudioConfig)
-                    
-                    if "error" not in result:
-                        st.session_state.audio_result = result
-                        st.session_state.audio_spec = spec_viz
-                        st.success("Analysis completed successfully!")
-                    else:
-                        st.error(f"Analysis failed: {result['error']}")
-        
-        if 'audio_result' in st.session_state:
-            result = st.session_state.audio_result
-            
-            st.markdown("---")
-            st.markdown("## DETECTION RESULTS")
-            
-            is_fake = result['predicted_class'] != "Original"
-            
-            if is_fake:
-                st.markdown(f"""
-                <div class="result-warning">
-                    <h2 class="result-title" style="color: #ef4444;">DEEPFAKE AUDIO DETECTED</h2>
-                    <p class="result-subtitle">Detected Synthesis Method: <strong>{result['predicted_class']}</strong></p>
-                    <p class="result-subtitle">Detection Confidence: <strong>{result['confidence']*100:.2f}%</strong></p>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                st.markdown(f"""
-                <div class="result-success">
-                    <h2 class="result-title" style="color: #10b981;">AUTHENTIC AUDIO</h2>
-                    <p class="result-subtitle">Classification: <strong>{result['predicted_class']}</strong></p>
-                    <p class="result-subtitle">Authenticity Confidence: <strong>{result['confidence']*100:.2f}%</strong></p>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            col1, col2, col3 = st.columns(3)
-            with col1:
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value">{result['predicted_class']}</div>
-                    <div class="metric-label">Predicted Class</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            with col2:
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value">{result['confidence']*100:.1f}%</div>
-                    <div class="metric-label">Confidence Score</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            with col3:
-                authenticity = "SYNTHETIC" if is_fake else "AUTHENTIC"
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value" style="color: {'#ef4444' if is_fake else '#10b981'};">{authenticity}</div>
-                    <div class="metric-label">Status</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            st.markdown("### Class Probability Distribution")
-            fig_probs = plot_probabilities(result['probabilities'], AudioConfig.CLASS_NAMES)
-            st.plotly_chart(fig_probs, use_container_width=True)
-            
-            if 'audio_spec' in st.session_state and st.session_state.audio_spec is not None:
-                st.markdown("### Mel Spectrogram Analysis")
-                fig_spec = plot_spectrogram(st.session_state.audio_spec)
-                st.plotly_chart(fig_spec, use_container_width=True)
-            
-            with st.expander("Detailed Class Probabilities"):
-                sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True)
-                for i, (class_name, prob) in enumerate(sorted_probs):
-                    st.markdown(f"**{i+1}. {class_name}:** `{prob*100:.4f}%`")
-
-def show_image_page():
-    st.markdown("""
-    <div class="gradient-header">
-        <h1>IMAGE DEEPFAKE DETECTION</h1>
-        <p>AI-Generated & Manipulated Image Analysis</p>
-    </div>
-    """, unsafe_allow_html=True)
-    
-    model = load_image_model()
-    if model is None:
-        st.error("Image model not loaded. Please ensure Image.pt is in the correct directory.")
-        return
-    
-    st.markdown("### Upload Image File for Analysis")
-    uploaded_file = st.file_uploader(
-        "Drag and drop your image file here or click to browse",
-        type=['jpg', 'jpeg', 'png', 'bmp', 'webp'],
-        help="Supported formats: JPG, PNG, BMP, WEBP"
-    )
-    
-    if uploaded_file is not None:
-        col1, col2 = st.columns([1, 1])
-        
-        with col1:
-            st.markdown("#### Image Preview")
-            image = Image.open(uploaded_file)
-            st.image(image, caption="Uploaded Image")
-            st.markdown(f"""
-            **File Information:**
-            - Filename: `{uploaded_file.name}`
-            - Size: `{uploaded_file.size / 1024:.2f} KB`
-            - Dimensions: `{image.size[0]} x {image.size[1]} pixels`
-            - Mode: `{image.mode}`
-            """)
-        
-        with col2:
-            st.markdown("#### Analysis Configuration")
-            st.markdown(f"""
-            **Model Settings:**
-            - Input Size: {ImageConfig.IMG_SIZE}x{ImageConfig.IMG_SIZE} pixels
-            - Architecture: Enhanced XceptionNet
-            - Device: {ImageConfig.DEVICE}
-            - Attention Mechanism: Enabled
-            """)
-            
-            if st.button("ANALYZE IMAGE", type="primary"):
-                with st.spinner("Processing image and extracting features..."):
-                    with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file:
-                        image.save(tmp_file.name)
-                        tmp_path = tmp_file.name
-                    
-                    result = predict_image(model, Path(tmp_path), ImageConfig)
-                    
-                    if "error" not in result:
-                        st.session_state.image_result = result
-                        st.success("Analysis completed successfully!")
-                    else:
-                        st.error(f"Analysis failed: {result['error']}")
-        
-        if 'image_result' in st.session_state:
-            result = st.session_state.image_result
-            
-            st.markdown("---")
-            st.markdown("## DETECTION RESULTS")
-            
-            is_fake = result['predicted_class'] != "Real"
-            
-            if is_fake:
-                st.markdown(f"""
-                <div class="result-warning">
-                    <h2 class="result-title" style="color: #ef4444;">SYNTHETIC IMAGE DETECTED</h2>
-                    <p class="result-subtitle">Generation Method: <strong>{result['predicted_class']}</strong></p>
-                    <p class="result-subtitle">Detection Confidence: <strong>{result['confidence']*100:.2f}%</strong></p>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                st.markdown(f"""
-                <div class="result-success">
-                    <h2 class="result-title" style="color: #10b981;">AUTHENTIC IMAGE</h2>
-                    <p class="result-subtitle">Classification: <strong>{result['predicted_class']}</strong></p>
-                    <p class="result-subtitle">Authenticity Confidence: <strong>{result['confidence']*100:.2f}%</strong></p>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            col1, col2, col3 = st.columns(3)
-            with col1:
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value">{result['predicted_class']}</div>
-                    <div class="metric-label">Predicted Class</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            with col2:
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value">{result['confidence']*100:.1f}%</div>
-                    <div class="metric-label">Confidence Score</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            with col3:
-                authenticity = "SYNTHETIC" if is_fake else "AUTHENTIC"
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value" style="color: {'#ef4444' if is_fake else '#10b981'};">{authenticity}</div>
-                    <div class="metric-label">Status</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            st.markdown("### Class Probability Distribution")
-            fig_probs = plot_probabilities(result['probabilities'], ImageConfig.CLASS_NAMES)
-            st.plotly_chart(fig_probs, use_container_width=True)
-            
-            st.markdown("### Prediction Confidence Matrix")
-            fig_cm = plot_confusion_matrix(result['probabilities'], result['predicted_class'], ImageConfig.CLASS_NAMES)
-            st.plotly_chart(fig_cm, use_container_width=True)
-            
-            with st.expander("Detailed Class Probabilities"):
-                sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True)
-                for i, (class_name, prob) in enumerate(sorted_probs):
-                    st.markdown(f"**{i+1}. {class_name}:** `{prob*100:.4f}%`")
-
-def show_video_page():
-    st.markdown("""
-    <div class="gradient-header">
-        <h1>VIDEO DEEPFAKE DETECTION</h1>
-        <p>Face Manipulation & Video Synthesis Analysis</p>
-    </div>
-    """, unsafe_allow_html=True)
-    
-    model = load_video_model()
-    if model is None:
-        st.error("Video model not loaded. Please ensure Video.pt is in the correct directory.")
-        return
-    
-    st.markdown("### Upload Video File for Analysis")
-    uploaded_file = st.file_uploader(
-        "Drag and drop your video file here or click to browse",
-        type=['mp4', 'avi', 'mov', 'mkv', 'webm'],
-        help="Supported formats: MP4, AVI, MOV, MKV, WEBM"
-    )
-    
-    if uploaded_file is not None:
-        col1, col2 = st.columns([1, 1])
-        
-        with col1:
-            st.markdown("#### Video Preview")
-            st.video(uploaded_file)
-            st.markdown(f"""
-            **File Information:**
-            - Filename: `{uploaded_file.name}`
-            - Size: `{uploaded_file.size / (1024*1024):.2f} MB`
-            - Type: `{uploaded_file.type}`
-            """)
-        
-        with col2:
-            st.markdown("#### Analysis Configuration")
-            frame_skip = st.slider(
-                "Frame Skip Rate", 
-                min_value=1, 
-                max_value=30, 
-                value=10,
-                help="Higher values = faster processing but may reduce accuracy"
-            )
-            
-            st.markdown(f"""
-            **Model Settings:**
-            - Architecture: Legacy Xception
-            - Frame Analysis: Every {frame_skip} frames
-            - Input Size: 300x300 pixels per frame
-            - Device: {VideoConfig.DEVICE}
-            """)
-            
-            if st.button("ANALYZE VIDEO", type="primary", use_container_width=True):
-                with st.spinner("Extracting frames and performing temporal analysis..."):
-                    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
-                        tmp_file.write(uploaded_file.getvalue())
-                        tmp_path = tmp_file.name
-                    
-                    VideoConfig.FRAME_SKIP = frame_skip
-                    result = predict_video(model, Path(tmp_path), VideoConfig)
-                    
-                    if "error" not in result:
-                        st.session_state.video_result = result
-                        st.success("Analysis completed successfully!")
-                    else:
-                        st.error(f"Analysis failed: {result['error']}")
-        
-        if 'video_result' in st.session_state:
-            result = st.session_state.video_result
-            
-            st.markdown("---")
-            st.markdown("## DETECTION RESULTS")
-            
-            is_fake = result['predicted_class'] != "Original"
-            
-            if is_fake:
-                st.markdown(f"""
-                <div class="result-warning">
-                    <h2 class="result-title" style="color: #ef4444;">DEEPFAKE VIDEO DETECTED</h2>
-                    <p class="result-subtitle">Manipulation Type: <strong>{result['predicted_class']}</strong></p>
-                    <p class="result-subtitle">Detection Confidence: <strong>{result['confidence']*100:.2f}%</strong></p>
-                    <p class="result-subtitle">Frames Analyzed: <strong>{result.get('frame_count', 'N/A')}</strong></p>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                st.markdown(f"""
-                <div class="result-success">
-                    <h2 class="result-title" style="color: #10b981;">AUTHENTIC VIDEO</h2>
-                    <p class="result-subtitle">Classification: <strong>{result['predicted_class']}</strong></p>
-                    <p class="result-subtitle">Authenticity Confidence: <strong>{result['confidence']*100:.2f}%</strong></p>
-                    <p class="result-subtitle">Frames Analyzed: <strong>{result.get('frame_count', 'N/A')}</strong></p>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            col1, col2, col3, col4 = st.columns(4)
-            with col1:
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value">{result['predicted_class']}</div>
-                    <div class="metric-label">Predicted Class</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            with col2:
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value">{result['confidence']*100:.1f}%</div>
-                    <div class="metric-label">Confidence</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            with col3:
-                authenticity = "SYNTHETIC" if is_fake else "AUTHENTIC"
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value" style="color: {'#ef4444' if is_fake else '#10b981'};">{authenticity}</div>
-                    <div class="metric-label">Status</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            with col4:
-                st.markdown(f"""
-                <div class="metric-card">
-                    <div class="metric-value">{result.get('frame_count', 0)}</div>
-                    <div class="metric-label">Frames</div>
-                </div>
-                """, unsafe_allow_html=True)
-            
-            st.markdown("### Class Probability Distribution")
-            fig_probs = plot_probabilities(result['probabilities'], VideoConfig.CLASS_NAMES)
-            st.plotly_chart(fig_probs, use_container_width=True)
-            
-            st.markdown("### Prediction Confidence Matrix")
-            fig_cm = plot_confusion_matrix(result['probabilities'], result['predicted_class'], VideoConfig.CLASS_NAMES)
-            st.plotly_chart(fig_cm, use_container_width=True)
-            
-            with st.expander("Detailed Class Probabilities"):
-                sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True)
-                for i, (class_name, prob) in enumerate(sorted_probs):
-                    st.markdown(f"**{i+1}. {class_name}:** `{prob*100:.4f}%`")
-
-# ======================= MAIN APP ============================
-def main():
-    with st.sidebar:
-        # st.markdown("""
-        # <div style="text-align: center; padding: 1.5rem; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 12px; margin-bottom: 2rem;">
-        #     <h2 style="color: white; margin: 0; font-size: 1.5rem; font-weight: 800;">DFDA SYSTEM</h2>
-        #     <p style="color: rgba(255,255,255,0.9); font-size: 0.95rem; margin: 0.5rem 0 0 0;">Deepfake Detection & Attribution</p>
-        # </div>
-        # """, unsafe_allow_html=True)
-        
-        page = st.radio(
-            "Navigation Menu",
-            ["Home", "Audio Detection", "Image Detection", "Video Detection"],
-            label_visibility="visible"
-        )
-        
-        st.markdown("---")
-        
-        st.markdown("""
-        <div style="padding: 1.2rem; background: rgba(102, 126, 234, 0.15); border-radius: 10px; border-left: 4px solid #667eea;">
-            <h4 style="margin-top: 0; color: #a5b4fc; font-size: 1.1rem;">System Status</h4>
-            <p style="margin: 0.5rem 0; color: #e4e4e7;"><strong>Device:</strong> {}</p>
-            <p style="margin: 0.5rem 0; color: #e4e4e7;"><strong>Models:</strong> Active</p>
-            <p style="margin: 0.5rem 0; color: #e4e4e7;"><strong>Status:</strong> Online</p>
-        </div>
-        """.format("GPU (CUDA)" if torch.cuda.is_available() else "CPU"), unsafe_allow_html=True)
-        
-        st.markdown("---")
-        
-        with st.expander("About DFDA System"):
-            st.markdown("""
-            **DFDA (Deepfake Detection & Attribution)** is an advanced AI-powered platform designed for detecting and attributing synthetic media content across multiple modalities.
-            
-            **Core Capabilities:**
-            - Multi-modal deepfake detection
-            - High-accuracy classification
-            - Real-time inference
-            - Comprehensive probability analysis
-            - Visual interpretability tools
-            
-            **Applications:**
-            - Academic research
-            - Media forensics
-            - Content verification
-            - Digital evidence analysis
-            
-            **Version:** 1.0.0  
-            **Framework:** PyTorch 2.0+
-            """)
-        
-        st.markdown("---")
-        st.markdown("""
-        <div style="text-align: left; color: #a5b4fc; font-size: 0.85rem; padding: 1rem 0;">
-            <p style="margin: 0;">Powered by Deep Learning
-                    Developed by:</p>
-            <p style="margin: 0.3rem 0 0 0;">
-                    2303A51L99 - Sri Varsha Janagam<br>
-                    2303A51LA0 - Sindhu Kodati<br>
-                    2303A51LA7 - Jayanth Bottu<br>
-                    2303A51LA9 - Mahandra Gaddam<br>
-                    2303A51LB0 - Nikhil Kuchana
-                    </p>
-        </div>
-        """, unsafe_allow_html=True)
-    
-    if page == "Home":
-        show_home_page()
-    elif page == "Audio Detection":
-        show_audio_page()
-    elif page == "Image Detection":
-        show_image_page()
-    elif page == "Video Detection":
-        show_video_page()
-
-if __name__ == "__main__":
-    main()
+# app.py — HF Spaces ready
+import streamlit as st
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+import torchaudio.transforms as T
+import torchvision.transforms as transforms
+import numpy as np
+import cv2
+from pathlib import Path
+from PIL import Image
+import plotly.graph_objects as go
+import plotly.express as px
+import io
+import timm
+import warnings
+import tempfile
+import shutil
+import os
+
+# ensure download_models.py is present next to app.py and exposes download_from_hf and MODELS_DIR
+from download_models import download_from_hf, MODELS_DIR
+
+warnings.filterwarnings('ignore')
+
+# HF repo & filenames (your repo)
+AUDIO_REPO = "jayanthbottu/DFDA"
+AUDIO_FILENAME = "Audio.pth"
+
+VIDEO_REPO = "jayanthbottu/DFDA"
+VIDEO_FILENAME = "Video.pth"
+
+IMAGE_REPO = "jayanthbottu/DFDA"
+IMAGE_FILENAME = "Image.pt"
+
+MODELS_DIR.mkdir(parents=True, exist_ok=True)
+
+# ======================= PAGE CONFIG ============================
+st.set_page_config(
+    page_title="DFDA - Deepfake Detection & Attribution",
+    page_icon="https://play-lh.googleusercontent.com/wJEqeRZGZds8tVZKQq-azZ1sp7IjYgF92zXVKYD1vy1LrV6xz4KVriABCjdjS8jQyA",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# ======================= CUSTOM CSS ============================
+st.markdown("""<style>/* ... keep your CSS unchanged ... */</style>""", unsafe_allow_html=True)
+
+# ======================= CONFIG CLASSES ============================
+class AudioConfig:
+    MODEL_PATH = MODELS_DIR / "Audio.pth"
+    CLASS_NAMES = ["ElevenLabs", "Original", "Tacotron", "Text To Speech", "Voice Conversion"]
+    SAMPLE_RATE = 16000
+    N_MELS = 128
+    N_FFT = 2048
+    HOP_LENGTH = 512
+    MAX_AUDIO_LENGTH = 5
+    DEVICE = torch.device("cpu")  # Spaces are CPU only
+
+class ImageConfig:
+    MODEL_PATH = MODELS_DIR / "Image.pt"
+    CLASS_NAMES = ["DALL-E", "DeepFaceLab", "Face2Face", "FaceShifter", "FaceSwap", "Midjourney",
+                   "NeuralTextures", "Real", "Stable Diffusion", "StyleGAN"]
+    IMG_SIZE = 299
+    DEVICE = torch.device("cpu")
+
+class VideoConfig:
+    MODEL_PATH = MODELS_DIR / "Video.pth"
+    CLASS_NAMES = ["Deepfakes", "Face2Face", "FaceShifter", "FaceSwap", "NeuralTextures", "Original"]
+    DEVICE = torch.device("cpu")
+    FRAME_SKIP = 10
+
+# ======================= MODEL DOWNLOAD & LOADING HELPERS ============================
+def download_and_load_torch_model(repo_id: str, filename: str, expected_device: torch.device):
+    """
+    Downloads filename from repo_id into MODELS_DIR (if missing) and returns local Path.
+    Raises exception on failure.
+    """
+    try:
+        local_path = download_from_hf(repo_id, filename)
+    except Exception as e:
+        raise RuntimeError(f"Download failed for {repo_id}/{filename}: {e}")
+    if not local_path.exists():
+        raise FileNotFoundError(f"After download, file not found: {local_path}")
+    # Loading deferred to specific loader functions (so we can instantiate architecture before loading state)
+    return local_path
+
+# ======================= AUDIO / IMAGE / VIDEO MODEL CLASSES (unchanged) ============================
+# Audio model class
+class AudioDeepfakeCNN(nn.Module):
+    def __init__(self, num_classes=5):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(1, 64, 3, padding=1),
+            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, 3, padding=1),
+            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
+            nn.MaxPool2d(2,2),
+            nn.Dropout2d(0.25)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(64, 128, 3, padding=1),
+            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
+            nn.MaxPool2d(2,2),
+            nn.Dropout2d(0.25)
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(128, 256, 3, padding=1),
+            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, 3, padding=1),
+            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
+            nn.MaxPool2d(2,2),
+            nn.Dropout2d(0.3)
+        )
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(256, 512, 3, padding=1),
+            nn.BatchNorm2d(512), nn.ReLU(inplace=True),
+            nn.Conv2d(512, 512, 3, padding=1),
+            nn.BatchNorm2d(512), nn.ReLU(inplace=True),
+            nn.MaxPool2d(2,2),
+            nn.Dropout2d(0.3)
+        )
+        self.gap = nn.AdaptiveAvgPool2d((1,1))
+        self.fc = nn.Sequential(
+            nn.Linear(512, 256), nn.ReLU(inplace=True), nn.Dropout(0.5),
+            nn.Linear(256, 128), nn.ReLU(inplace=True), nn.Dropout(0.4),
+            nn.Linear(128, num_classes)
+        )
+    def forward(self, x):
+        x = self.conv1(x); x = self.conv2(x); x = self.conv3(x); x = self.conv4(x)
+        x = self.gap(x); x = x.view(x.size(0), -1); x = self.fc(x); return x
+
+# Image model pieces
+class AttentionBlock(nn.Module):
+    def __init__(self, in_features):
+        super(AttentionBlock, self).__init__()
+        self.attention = nn.Sequential(
+            nn.Linear(in_features, in_features // 8),
+            nn.ReLU(),
+            nn.Linear(in_features // 8, in_features),
+            nn.Sigmoid()
+        )
+    def forward(self, x):
+        att_weights = self.attention(x)
+        return x * att_weights
+
+class EnhancedXceptionNet(nn.Module):
+    def __init__(self, num_classes, dropout_rate=0.3, use_attention=True):
+        super(EnhancedXceptionNet, self).__init__()
+        self.base_model = timm.create_model('xception', pretrained=False)
+        in_features = self.base_model.get_classifier().in_features
+        self.base_model.reset_classifier(0)
+        self.use_attention = use_attention
+        if use_attention:
+            self.attention = AttentionBlock(in_features)
+        self.classifier = nn.Sequential(
+            nn.Dropout(dropout_rate),
+            nn.Linear(in_features, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate / 2),
+            nn.Linear(512, num_classes)
+        )
+    def forward(self, x):
+        features = self.base_model(x)
+        if self.use_attention:
+            features = self.attention(features)
+        output = self.classifier(features)
+        return output
+
+# ======================= MODEL LOADERS (use cache_resource for HF Spaces) ============================
+@st.cache_resource
+def load_audio_model():
+    # instantiate and then load state_dict
+    try:
+        model = AudioDeepfakeCNN(num_classes=len(AudioConfig.CLASS_NAMES))
+        checkpoint = torch.load(AudioConfig.MODEL_PATH, map_location="cpu")
+        state = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
+        model.load_state_dict(state, strict=False)
+        model.to(AudioConfig.DEVICE).eval()
+        return model
+    except Exception as e:
+        st.error(f"Error loading audio model: {e}")
+        return None
+
+@st.cache_resource
+def load_image_model():
+    try:
+        model = EnhancedXceptionNet(num_classes=len(ImageConfig.CLASS_NAMES), dropout_rate=0.3, use_attention=True)
+        checkpoint = torch.load(ImageConfig.MODEL_PATH, map_location="cpu")
+        state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
+        model_dict = model.state_dict()
+        filtered_dict = {k: v for k, v in state_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
+        model_dict.update(filtered_dict)
+        model.load_state_dict(model_dict, strict=False)
+        model.to(ImageConfig.DEVICE).eval()
+        return model
+    except Exception as e:
+        st.error(f"Error loading image model: {e}")
+        return None
+
+@st.cache_resource
+def load_video_model():
+    try:
+        num_classes = len(VideoConfig.CLASS_NAMES)
+        model = timm.create_model("legacy_xception", pretrained=False, num_classes=num_classes)
+        checkpoint = torch.load(VideoConfig.MODEL_PATH, map_location="cpu")
+        state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
+        model_dict = model.state_dict()
+        filtered_dict = {k: v for k, v in state_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
+        model_dict.update(filtered_dict)
+        model.load_state_dict(model_dict, strict=False)
+        model.to(VideoConfig.DEVICE).eval()
+        return model
+    except Exception as e:
+        st.error(f"Error loading video model: {e}")
+        return None
+
+# ======================= PREPROCESSING (unchanged) ============================
+def preprocess_audio(audio_path, sample_rate, n_mels, n_fft, hop_length, max_length):
+    waveform, sr = torchaudio.load(str(audio_path))
+    if waveform.shape[0] > 1:
+        waveform = torch.mean(waveform, dim=0, keepdim=True)
+    if sr != sample_rate:
+        resampler = T.Resample(sr, sample_rate)
+        waveform = resampler(waveform)
+    max_len = sample_rate * max_length
+    if waveform.shape[1] < max_len:
+        waveform = torch.nn.functional.pad(waveform, (0, max_len - waveform.shape[1]))
+    else:
+        waveform = waveform[:, :max_len]
+    mel_spec = T.MelSpectrogram(sample_rate, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels, power=2.0)(waveform)
+    mel_spec_db = T.AmplitudeToDB(top_db=80)(mel_spec)
+    mel_spec_db = (mel_spec_db - mel_spec_db.mean()) / (mel_spec_db.std() + 1e-8)
+    return mel_spec_db, mel_spec_db.squeeze().numpy()
+
+def preprocess_image(image_path, size):
+    trans = transforms.Compose([
+        transforms.Resize((size, size)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+    ])
+    img = Image.open(image_path).convert('RGB')
+    return trans(img).unsqueeze(0)
+
+def extract_frames(video_path, frame_skip=10):
+    frames = []
+    cap = cv2.VideoCapture(str(video_path))
+    count = 0
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret: break
+        if count % frame_skip == 0:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frames.append(frame)
+        count += 1
+    cap.release()
+    return frames
+
+def preprocess_frame(frame, size):
+    trans = transforms.Compose([
+        transforms.ToPILImage(),
+        transforms.Resize((size, size)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+    ])
+    return trans(frame).unsqueeze(0)
+
+# ======================= INFERENCE (unchanged) ============================
+def predict_audio(model, audio_path, config):
+    device = config.DEVICE
+    try:
+        x, spec_viz = preprocess_audio(audio_path, config.SAMPLE_RATE, config.N_MELS, config.N_FFT, config.HOP_LENGTH, config.MAX_AUDIO_LENGTH)
+        x = x.unsqueeze(0).to(device)
+        with torch.no_grad():
+            out = F.softmax(model(x), dim=1)[0]
+        pred_idx = torch.argmax(out).item()
+        return {
+            "predicted_class": config.CLASS_NAMES[pred_idx],
+            "predicted_index": pred_idx,
+            "confidence": out[pred_idx].item(),
+            "probabilities": {c: out[i].item() for i, c in enumerate(config.CLASS_NAMES)}
+        }, spec_viz
+    except Exception as e:
+        return {"error": str(e)}, None
+
+def predict_image(model, image_path, config):
+    device = config.DEVICE
+    try:
+        x = preprocess_image(image_path, config.IMG_SIZE).to(device)
+        with torch.no_grad():
+            out = F.softmax(model(x), dim=1)[0]
+        pred_idx = torch.argmax(out).item()
+        return {
+            "predicted_class": config.CLASS_NAMES[pred_idx],
+            "predicted_index": pred_idx,
+            "confidence": out[pred_idx].item(),
+            "probabilities": {c: out[i].item() for i, c in enumerate(config.CLASS_NAMES)}
+        }
+    except Exception as e:
+        return {"error": str(e)}
+
+def predict_video(model, video_path, config):
+    device = config.DEVICE
+    try:
+        frames = extract_frames(video_path, config.FRAME_SKIP)
+        if len(frames) == 0:
+            return {"error":"No frames extracted"}
+        probs = []
+        progress_bar = st.progress(0)
+        for i, frame in enumerate(frames):
+            x = preprocess_frame(frame, 300).to(device)
+            with torch.no_grad():
+                p = F.softmax(model(x),dim=1)[0].cpu().numpy()
+            probs.append(p)
+            progress_bar.progress((i + 1) / len(frames))
+        probs = np.stack(probs)
+        avg_prob = probs.mean(axis=0)
+        pred_idx = np.argmax(avg_prob)
+        return {
+            "predicted_class": config.CLASS_NAMES[pred_idx],
+            "predicted_index": pred_idx,
+            "confidence": float(avg_prob[pred_idx]),
+            "probabilities": {c: float(avg_prob[i]) for i, c in enumerate(config.CLASS_NAMES)},
+            "frame_count": len(frames)
+        }
+    except Exception as e:
+        return {"error": str(e)}
+
+# ======================= VISUALIZATION (unchanged) ============================
+def plot_probabilities(probs, class_names):
+    fig = go.Figure(data=[
+        go.Bar(
+            x=list(probs.values()),
+            y=list(probs.keys()),
+            orientation='h',
+            marker=dict(
+                color=list(probs.values()),
+                colorscale=[[0, '#667eea'], [0.5, '#764ba2'], [1, '#f093fb']],
+                line=dict(color='rgba(102, 126, 234, 0.8)', width=2)
+            ),
+            text=[f'{v*100:.2f}%' for v in probs.values()],
+            textposition='outside',
+            textfont=dict(size=14, color='#e4e4e7', family='Inter')
+        )
+    ])
+    fig.update_layout(
+        title=dict(text="Class Probability Distribution", font=dict(size=22, color='#a5b4fc', family='Inter')),
+        xaxis_title="Probability",
+        yaxis_title="Class",
+        height=400,
+        template="plotly_dark",
+        showlegend=False,
+        plot_bgcolor='rgba(15,15,30,0.5)',
+        paper_bgcolor='rgba(0,0,0,0)',
+        font=dict(color='#e4e4e7', family='Inter'),
+        xaxis=dict(gridcolor='rgba(102, 126, 234, 0.2)'),
+        yaxis=dict(gridcolor='rgba(102, 126, 234, 0.2)')
+    )
+    return fig
+
+def plot_spectrogram(spec_data):
+    if spec_data is None:
+        return None
+    fig = px.imshow(
+        spec_data,
+        labels=dict(x="Time", y="Frequency", color="Amplitude (dB)"),
+        aspect='auto',
+        color_continuous_scale=[[0, '#0f0f1e'], [0.25, '#667eea'], [0.5, '#764ba2'], [0.75, '#f093fb'], [1, '#fbbf24']]
+    )
+    fig.update_layout(
+        title=dict(text="Mel Spectrogram Visualization", font=dict(size=22, color='#a5b4fc', family='Inter')),
+        height=400,
+        template="plotly_dark",
+        plot_bgcolor='rgba(15,15,30,0.5)',
+        paper_bgcolor='rgba(0,0,0,0)',
+        font=dict(color='#e4e4e7', family='Inter')
+    )
+    return fig
+
+def plot_confusion_matrix(probs, predicted_class, class_names):
+    n_classes = len(class_names)
+    cm = np.zeros((n_classes, n_classes))
+    pred_idx = class_names.index(predicted_class)
+    for i, (cls, prob) in enumerate(probs.items()):
+        cm[pred_idx, i] = prob
+    fig = px.imshow(
+        cm,
+        labels=dict(x="Predicted Class", y="True Class", color="Probability"),
+        x=class_names,
+        y=class_names,
+        color_continuous_scale=[[0, '#0f0f1e'], [0.5, '#667eea'], [1, '#f093fb']],
+        aspect='auto',
+        text_auto='.2f'
+    )
+    fig.update_layout(
+        title=dict(text="Prediction Confidence Matrix", font=dict(size=22, color='#a5b4fc', family='Inter')),
+        height=500,
+        template="plotly_dark",
+        plot_bgcolor='rgba(15,15,30,0.5)',
+        paper_bgcolor='rgba(0,0,0,0)',
+        font=dict(color='#e4e4e7', family='Inter')
+    )
+    return fig
+
+# ======================= STARTUP: download and set model paths ============================
+# This runs on import/startup — keeps UI responsive by using messages
+with st.spinner("Downloading models from Hugging Face (cached to /app/models)... This will happen once per Space runtime."):
+    try:
+        audio_path = download_and_load_torch_model(AUDIO_REPO, AUDIO_FILENAME, AudioConfig.DEVICE)
+        AudioConfig.MODEL_PATH = audio_path
+    except Exception as e:
+        st.error(f"Audio model download/load error: {e}")
+        # do not stop app — keep UI showing error
+    try:
+        image_path = download_and_load_torch_model(IMAGE_REPO, IMAGE_FILENAME, ImageConfig.DEVICE)
+        ImageConfig.MODEL_PATH = image_path
+    except Exception as e:
+        st.error(f"Image model download/load error: {e}")
+    try:
+        video_path = download_and_load_torch_model(VIDEO_REPO, VIDEO_FILENAME, VideoConfig.DEVICE)
+        VideoConfig.MODEL_PATH = video_path
+    except Exception as e:
+        st.error(f"Video model download/load error: {e}")
+
+# Load models (cached resource)
+audio_model = load_audio_model()
+image_model = load_image_model()
+video_model = load_video_model()
+
+# ======================= PAGES ============================
+# (keep all your page functions unchanged: show_home_page, show_audio_page, show_image_page, show_video_page)
+# ... paste the rest of your existing functions here (they already reference load_* and configs) ...
+
+# For brevity in this listing, we assume the rest of your UI functions are unchanged and present
+# (You can reuse the exact definitions of show_home_page, show_audio_page, show_image_page, show_video_page
+#  from your original file — they will work as-is because AudioConfig, ImageConfig, VideoConfig were updated.)
+
+def main():
+    with st.sidebar:
+        page = st.radio(
+            "Navigation Menu",
+            ["Home", "Audio Detection", "Image Detection", "Video Detection"],
+            label_visibility="visible"
+        )
+        st.markdown("---")
+        st.markdown(f"""
+        <div style="padding: 1.2rem; background: rgba(102, 126, 234, 0.15); border-radius: 10px; border-left: 4px solid #667eea;">
+            <h4 style="margin-top: 0; color: #a5b4fc; font-size: 1.1rem;">System Status</h4>
+            <p style="margin: 0.5rem 0; color: #e4e4e7;"><strong>Device:</strong> {"CPU"}</p>
+            <p style="margin: 0.5rem 0; color: #e4e4e7;"><strong>Models:</strong> Active</p>
+            <p style="margin: 0.5rem 0; color: #e4e4e7;"><strong>Status:</strong> Online</p>
+        </div>
+        """, unsafe_allow_html=True)
+
+    if st.session_state.get("initialized") is None:
+        st.session_state["initialized"] = True
+
+    # route pages
+    page = st.sidebar.radio("Navigation Menu", ["Home", "Audio Detection", "Image Detection", "Video Detection"])
+    if page == "Home":
+        show_home_page()
+    elif page == "Audio Detection":
+        show_audio_page()
+    elif page == "Image Detection":
+        show_image_page()
+    elif page == "Video Detection":
+        show_video_page()
+
+if __name__ == "__main__":
+    main()