app.py

import os
from io import BytesIO
import uuid
import threading

import cv2
import numpy as np
import torch

from fastapi import FastAPI, UploadFile, File
from fastapi.responses import HTMLResponse, StreamingResponse, FileResponse

from huggingface_hub import hf_hub_download
from torchvision.transforms.functional import normalize

from facelib.utils.face_restoration_helper import FaceRestoreHelper
from models import CodeFormer
from utils import img2tensor, tensor2img

app = FastAPI(
    title="CodeFormer Face Enhancement",
    description="Face Restoration API",
    version="1.0.0"
)

# ====================================================
# Job Storage
# ====================================================

jobs = {}
results = {}

def update_progress(job_id, value, message):
    jobs[job_id] = {
        "progress": value,
        "message": message
    }

# ====================================================
# Load Model Once
# ====================================================

REPO_ID = "leonelhs/gfpgan"

pretrain_model_path = hf_hub_download(
    repo_id=REPO_ID,
    filename="CodeFormer.pth"
)

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

net = CodeFormer(
    dim_embd=512,
    codebook_size=1024,
    n_head=8,
    n_layers=9,
    connect_list=["32", "64", "128", "256"]
).to(device)

checkpoint = torch.load(
    pretrain_model_path,
    map_location=device
)["params_ema"]

net.load_state_dict(checkpoint)
net.eval()

face_helper = FaceRestoreHelper(
    upscale_factor=2,
    face_size=512,
    crop_ratio=(1, 1),
    det_model="retinaface_resnet50",
    save_ext="png",
    use_parse=True,
    device=device
)

# ====================================================
# Face Enhancement Function
# ====================================================

def enhance_face(image_rgb, job_id=None):
    
    face_helper.clean_all()
    
    if job_id:
        update_progress(job_id, 5, "Preparing image")
    
    image_bgr = cv2.cvtColor(
        image_rgb,
        cv2.COLOR_RGB2BGR
    )

    face_helper.read_image(image_bgr)

    if job_id:
        update_progress(job_id, 20, "Detecting faces")
    
    face_helper.get_face_landmarks_5(
        only_center_face=False,
        resize=640,
        eye_dist_threshold=5
    )

    if job_id:
        update_progress(job_id, 35, "Aligning faces")
    
    face_helper.align_warp_face()
    
    total_faces = len(face_helper.cropped_faces)
    
    if total_faces == 0:
        if job_id:
            update_progress(job_id, 100, "No faces found")
        return image_rgb

    for idx, cropped_face in enumerate(face_helper.cropped_faces):
        
        if job_id:
            progress = 40 + int((idx / total_faces) * 50)
            update_progress(
                job_id,
                progress,
                f"Enhancing face {idx+1}/{total_faces}"
            )

        cropped_face_t = img2tensor(
            cropped_face / 255.0,
            bgr2rgb=True,
            float32=True
        )

        normalize(
            cropped_face_t,
            (0.5, 0.5, 0.5),
            (0.5, 0.5, 0.5),
            inplace=True
        )

        cropped_face_t = cropped_face_t.unsqueeze(0).to(device)

        try:

            with torch.no_grad():

                output = net(
                    cropped_face_t,
                    w=0.5,
                    adain=True
                )[0]

                restored_face = tensor2img(
                    output,
                    rgb2bgr=True,
                    min_max=(-1, 1)
                )

            restored_face = restored_face.astype("uint8")

        except Exception as e:

            print(e)

            restored_face = tensor2img(
                cropped_face_t,
                rgb2bgr=True,
                min_max=(-1, 1)
            )

        face_helper.add_restored_face(
            restored_face,
            cropped_face
        )

    if job_id:
        update_progress(job_id, 95, "Finalizing image")
    
    face_helper.get_inverse_affine(None)

    restored_img = face_helper.paste_faces_to_input_image()

    restored_img = cv2.cvtColor(
        restored_img,
        cv2.COLOR_BGR2RGB
    )
    
    if job_id:
        update_progress(job_id, 100, "Completed")

    return restored_img


# ====================================================
# Background Worker
# ====================================================

def process_job(job_id, image):
    try:
        result = enhance_face(image, job_id)
        
        result = cv2.cvtColor(
            result,
            cv2.COLOR_RGB2BGR
        )
        
        _, buffer = cv2.imencode(
            ".png",
            result
        )
        
        results[job_id] = buffer.tobytes()
        
        update_progress(
            job_id,
            100,
            "Completed"
        )
        
    except Exception as e:
        update_progress(
            job_id,
            -1,
            str(e)
        )

# ====================================================
# UI
# ====================================================

HTML_PAGE = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>CodeFormer Face Enhancement</title>
    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap" rel="stylesheet">
    <style>
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }

        body {
            font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
            display: flex;
            justify-content: center;
            align-items: center;
        }

        .main-container {
            max-width: 1200px;
            width: 100%;
            margin: 0 auto;
        }

        .card {
            background: rgba(255, 255, 255, 0.95);
            backdrop-filter: blur(20px);
            border-radius: 24px;
            padding: 40px;
            box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
            border: 1px solid rgba(255, 255, 255, 0.2);
        }

        h1 {
            text-align: center;
            font-size: 2.5rem;
            font-weight: 700;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            -webkit-background-clip: text;
            -webkit-text-fill-color: transparent;
            background-clip: text;
            margin-bottom: 8px;
        }

        .subtitle {
            text-align: center;
            color: #6b7280;
            font-size: 1rem;
            margin-bottom: 30px;
            font-weight: 400;
        }

        .upload-area {
            border: 2px dashed #d1d5db;
            border-radius: 16px;
            padding: 40px;
            text-align: center;
            cursor: pointer;
            transition: all 0.3s ease;
            background: #f9fafb;
            position: relative;
        }

        .upload-area:hover {
            border-color: #667eea;
            background: #f3f4f6;
            transform: translateY(-2px);
        }

        .upload-area.dragover {
            border-color: #667eea;
            background: #eef2ff;
            transform: scale(1.02);
        }

        .upload-icon {
            font-size: 48px;
            margin-bottom: 12px;
        }

        .upload-text {
            color: #374151;
            font-size: 1.1rem;
            font-weight: 500;
        }

        .upload-subtext {
            color: #6b7280;
            font-size: 0.9rem;
            margin-top: 4px;
        }

        #file {
            display: none;
        }

        .controls {
            display: flex;
            gap: 12px;
            margin-top: 20px;
            flex-wrap: wrap;
        }

        .btn {
            padding: 12px 32px;
            border: none;
            border-radius: 12px;
            font-size: 1rem;
            font-weight: 600;
            cursor: pointer;
            transition: all 0.3s ease;
            display: inline-flex;
            align-items: center;
            gap: 8px;
            font-family: 'Inter', sans-serif;
        }

        .btn-primary {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4);
        }

        .btn-primary:hover:not(:disabled) {
            transform: translateY(-2px);
            box-shadow: 0 6px 20px rgba(102, 126, 234, 0.5);
        }

        .btn-primary:disabled {
            opacity: 0.5;
            cursor: not-allowed;
            transform: none;
        }

        .btn-success {
            background: #10b981;
            color: white;
            box-shadow: 0 4px 15px rgba(16, 185, 129, 0.4);
        }

        .btn-success:hover:not(:disabled) {
            transform: translateY(-2px);
            box-shadow: 0 6px 20px rgba(16, 185, 129, 0.5);
        }

        .btn-success:disabled {
            opacity: 0.5;
            cursor: not-allowed;
            transform: none;
        }

        .btn-secondary {
            background: #f3f4f6;
            color: #374151;
        }

        .btn-secondary:hover {
            background: #e5e7eb;
            transform: translateY(-2px);
        }

        .btn-danger {
            background: #ef4444;
            color: white;
        }

        .btn-danger:hover {
            background: #dc2626;
            transform: translateY(-2px);
        }

        .status-container {
            margin-top: 20px;
            padding: 16px;
            border-radius: 12px;
            background: #f9fafb;
            display: none;
        }

        .status-container.active {
            display: block;
            animation: slideDown 0.3s ease;
        }

        @keyframes slideDown {
            from {
                opacity: 0;
                transform: translateY(-10px);
            }
            to {
                opacity: 1;
                transform: translateY(0);
            }
        }

        .status-header {
            display: flex;
            justify-content: space-between;
            align-items: center;
            margin-bottom: 8px;
        }

        .status-message {
            color: #374151;
            font-weight: 500;
            font-size: 0.95rem;
        }

        .status-percentage {
            color: #667eea;
            font-weight: 700;
            font-size: 1.1rem;
        }

        .progress-wrapper {
            width: 100%;
            height: 8px;
            background: #e5e7eb;
            border-radius: 100px;
            overflow: hidden;
            position: relative;
        }

        .progress-bar {
            height: 100%;
            width: 0%;
            background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
            border-radius: 100px;
            transition: width 0.5s cubic-bezier(0.4, 0, 0.2, 1);
            position: relative;
        }

        .progress-bar::after {
            content: '';
            position: absolute;
            top: 0;
            left: 0;
            right: 0;
            bottom: 0;
            background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.3), transparent);
            animation: shimmer 2s infinite;
        }

        @keyframes shimmer {
            0% {
                transform: translateX(-100%);
            }
            100% {
                transform: translateX(100%);
            }
        }

        .progress-bar.complete::after {
            animation: none;
        }

        .image-grid {
            display: grid;
            grid-template-columns: 1fr 1fr;
            gap: 24px;
            margin-top: 24px;
        }

        @media (max-width: 768px) {
            .image-grid {
                grid-template-columns: 1fr;
                gap: 16px;
            }
        }

        .image-card {
            background: #f9fafb;
            border-radius: 16px;
            padding: 16px;
            transition: all 0.3s ease;
        }

        .image-card:hover {
            transform: translateY(-4px);
            box-shadow: 0 8px 25px rgba(0, 0, 0, 0.1);
        }

        .image-label {
            font-weight: 600;
            color: #374151;
            margin-bottom: 12px;
            display: flex;
            align-items: center;
            gap: 8px;
            font-size: 0.95rem;
        }

        .image-label .badge {
            background: #667eea;
            color: white;
            padding: 2px 10px;
            border-radius: 100px;
            font-size: 0.75rem;
            font-weight: 600;
        }

        .image-label .badge-success {
            background: #10b981;
        }

        .image-container {
            width: 100%;
            aspect-ratio: 1;
            background: #e5e7eb;
            border-radius: 12px;
            overflow: hidden;
            position: relative;
            display: flex;
            align-items: center;
            justify-content: center;
        }

        .image-container img {
            width: 100%;
            height: 100%;
            object-fit: cover;
            transition: opacity 0.5s ease;
        }

        .image-container .placeholder {
            color: #9ca3af;
            font-size: 0.9rem;
            text-align: center;
            padding: 20px;
        }

        .image-container .placeholder .icon {
            font-size: 32px;
            display: block;
            margin-bottom: 8px;
        }

        .image-container.loading::after {
            content: '';
            position: absolute;
            top: 0;
            left: 0;
            right: 0;
            bottom: 0;
            background: linear-gradient(90deg, #e5e7eb, #f3f4f6, #e5e7eb);
            background-size: 200% 100%;
            animation: loading 1.5s infinite;
        }

        @keyframes loading {
            0% {
                background-position: -200% 0;
            }
            100% {
                background-position: 200% 0;
            }
        }

        .file-info {
            display: flex;
            align-items: center;
            gap: 12px;
            padding: 12px 16px;
            background: #f3f4f6;
            border-radius: 8px;
            margin-top: 12px;
            font-size: 0.9rem;
            color: #374151;
        }

        .file-info .file-name {
            font-weight: 500;
            flex: 1;
            overflow: hidden;
            text-overflow: ellipsis;
            white-space: nowrap;
        }

        .file-info .file-size {
            color: #6b7280;
            font-size: 0.85rem;
        }

        .file-info .remove-file {
            cursor: pointer;
            color: #ef4444;
            font-weight: 600;
            padding: 4px 8px;
            border-radius: 4px;
            transition: background 0.2s;
        }

        .file-info .remove-file:hover {
            background: #fee2e2;
        }

        .toast {
            position: fixed;
            top: 20px;
            right: 20px;
            padding: 16px 24px;
            border-radius: 12px;
            color: white;
            font-weight: 500;
            box-shadow: 0 10px 40px rgba(0, 0, 0, 0.2);
            transform: translateX(400px);
            transition: transform 0.3s ease;
            z-index: 1000;
            max-width: 400px;
        }

        .toast.show {
            transform: translateX(0);
        }

        .toast.success {
            background: #10b981;
        }

        .toast.error {
            background: #ef4444;
        }

        .toast.info {
            background: #3b82f6;
        }

        .download-section {
            display: none;
            margin-top: 16px;
            padding: 16px;
            background: #f0fdf4;
            border-radius: 12px;
            border: 1px solid #86efac;
            align-items: center;
            justify-content: space-between;
            flex-wrap: wrap;
            gap: 12px;
        }

        .download-section.active {
            display: flex;
            animation: slideDown 0.3s ease;
        }

        .download-info {
            display: flex;
            align-items: center;
            gap: 12px;
            color: #065f46;
        }

        .download-info .icon {
            font-size: 24px;
        }

        @media (max-width: 640px) {
            .card {
                padding: 20px;
            }

            h1 {
                font-size: 1.8rem;
            }

            .controls {
                flex-direction: column;
            }

            .btn {
                width: 100%;
                justify-content: center;
            }

            .upload-area {
                padding: 24px;
            }

            .toast {
                left: 20px;
                right: 20px;
                max-width: none;
            }

            .download-section {
                flex-direction: column;
                align-items: stretch;
            }
        }
    </style>
</head>
<body>

    <div class="main-container">
        <div class="card">
            <h1>✨ Face Enhancer</h1>
            <p class="subtitle">AI-powered face restoration using CodeFormer</p>

            <!-- Upload Area -->
            <div class="upload-area" id="uploadArea">
                <div class="upload-icon">📸</div>
                <div class="upload-text">Drop your image here or click to browse</div>
                <div class="upload-subtext">Supports JPG, PNG, WEBP (Max 10MB)</div>
                <input type="file" id="file" accept="image/*">
            </div>

            <!-- File Info -->
            <div id="fileInfo" class="file-info" style="display:none;">
                <span>📎</span>
                <span class="file-name" id="fileName">image.jpg</span>
                <span class="file-size" id="fileSize">2.4 MB</span>
                <span class="remove-file" id="removeFile">✕</span>
            </div>

            <!-- Controls -->
            <div class="controls">
                <button class="btn btn-primary" id="enhanceBtn" disabled>
                    🚀 Enhance Face
                </button>
                <button class="btn btn-secondary" id="resetBtn">
                    🔄 Reset
                </button>
            </div>

            <!-- Status -->
            <div class="status-container" id="statusContainer">
                <div class="status-header">
                    <span class="status-message" id="statusMessage">Processing...</span>
                    <span class="status-percentage" id="statusPercentage">0%</span>
                </div>
                <div class="progress-wrapper">
                    <div class="progress-bar" id="progressBar"></div>
                </div>
            </div>

            <!-- Download Section -->
            <div class="download-section" id="downloadSection">
                <div class="download-info">
                    <span class="icon">✅</span>
                    <span>Enhancement complete! Download your image</span>
                </div>
                <button class="btn btn-success" id="downloadBtn">
                    ⬇️ Download PNG
                </button>
            </div>

            <!-- Image Grid -->
            <div class="image-grid">
                <div class="image-card">
                    <div class="image-label">
                        <span>📷 Original</span>
                        <span class="badge">Input</span>
                    </div>
                    <div class="image-container" id="inputContainer">
                        <div class="placeholder">
                            <span class="icon">🖼️</span>
                            No image selected
                        </div>
                        <img id="inputPreview" style="display:none;" alt="Original">
                    </div>
                </div>

                <div class="image-card">
                    <div class="image-label">
                        <span>✨ Enhanced</span>
                        <span class="badge badge-success">Output</span>
                    </div>
                    <div class="image-container" id="outputContainer">
                        <div class="placeholder">
                            <span class="icon">🎯</span>
                            Waiting for enhancement
                        </div>
                        <img id="outputPreview" style="display:none;" alt="Enhanced">
                    </div>
                </div>
            </div>
        </div>
    </div>

    <!-- Toast Notification -->
    <div class="toast" id="toast"></div>

    <script>
        // DOM Elements
        const fileInput = document.getElementById('file');
        const uploadArea = document.getElementById('uploadArea');
        const fileInfo = document.getElementById('fileInfo');
        const fileName = document.getElementById('fileName');
        const fileSize = document.getElementById('fileSize');
        const removeFile = document.getElementById('removeFile');
        const enhanceBtn = document.getElementById('enhanceBtn');
        const resetBtn = document.getElementById('resetBtn');
        const downloadBtn = document.getElementById('downloadBtn');
        const downloadSection = document.getElementById('downloadSection');
        const inputPreview = document.getElementById('inputPreview');
        const outputPreview = document.getElementById('outputPreview');
        const inputContainer = document.getElementById('inputContainer');
        const outputContainer = document.getElementById('outputContainer');
        const statusContainer = document.getElementById('statusContainer');
        const statusMessage = document.getElementById('statusMessage');
        const statusPercentage = document.getElementById('statusPercentage');
        const progressBar = document.getElementById('progressBar');
        const toast = document.getElementById('toast');

        let selectedFile = null;
        let pollingInterval = null;
        let currentJobId = null;
        let currentResultUrl = null;

        // Toast function
        function showToast(message, type = 'info') {
            toast.textContent = message;
            toast.className = `toast ${type}`;
            setTimeout(() => {
                toast.classList.add('show');
            }, 10);
            setTimeout(() => {
                toast.classList.remove('show');
            }, 3000);
        }

        // Format file size
        function formatFileSize(bytes) {
            if (bytes < 1024) return bytes + ' B';
            if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB';
            return (bytes / (1024 * 1024)).toFixed(1) + ' MB';
        }

        // Update UI with selected file
        function handleFileSelect(file) {
            if (!file) return;

            // Validate file type
            if (!file.type.startsWith('image/')) {
                showToast('Please select an image file', 'error');
                return;
            }

            // Validate file size (10MB)
            if (file.size > 10 * 1024 * 1024) {
                showToast('File size must be less than 10MB', 'error');
                return;
            }

            selectedFile = file;
            fileName.textContent = file.name;
            fileSize.textContent = formatFileSize(file.size);
            fileInfo.style.display = 'flex';
            enhanceBtn.disabled = false;

            // Show preview
            const reader = new FileReader();
            reader.onload = function(e) {
                inputPreview.src = e.target.result;
                inputPreview.style.display = 'block';
                inputContainer.querySelector('.placeholder').style.display = 'none';
            };
            reader.readAsDataURL(file);

            showToast(`Loaded: ${file.name}`, 'success');
        }

        // Remove file
        function removeSelectedFile() {
            selectedFile = null;
            fileInput.value = '';
            fileInfo.style.display = 'none';
            enhanceBtn.disabled = true;
            inputPreview.style.display = 'none';
            inputContainer.querySelector('.placeholder').style.display = 'block';
            outputPreview.style.display = 'none';
            outputContainer.querySelector('.placeholder').style.display = 'block';
            downloadSection.classList.remove('active');
            if (currentResultUrl) {
                URL.revokeObjectURL(currentResultUrl);
                currentResultUrl = null;
            }
            resetProgress();
            if (pollingInterval) {
                clearInterval(pollingInterval);
                pollingInterval = null;
            }
        }

        // Reset progress UI
        function resetProgress() {
            statusContainer.classList.remove('active');
            statusMessage.textContent = 'Processing...';
            statusPercentage.textContent = '0%';
            progressBar.style.width = '0%';
            progressBar.classList.remove('complete');
        }

        // Update progress
        function updateProgress(progress) {
            statusContainer.classList.add('active');
            statusMessage.textContent = progress.message || 'Processing...';
            const value = Math.min(Math.max(progress.progress, 0), 100);
            statusPercentage.textContent = Math.round(value) + '%';
            progressBar.style.width = value + '%';
            
            if (value >= 100) {
                progressBar.classList.add('complete');
            }
        }

        // Download function
        function downloadImage() {
            if (!currentResultUrl) {
                showToast('No image to download', 'error');
                return;
            }
            
            // Create a temporary link element
            const link = document.createElement('a');
            link.href = currentResultUrl;
            
            // Generate filename from original file
            let baseName = selectedFile ? selectedFile.name.replace(/\.[^/.]+$/, '') : 'enhanced';
            link.download = `${baseName}_enhanced.png`;
            
            document.body.appendChild(link);
            link.click();
            document.body.removeChild(link);
            
            showToast('📥 Download started!', 'success');
        }

        // Enhance function
        async function enhance() {
            if (!selectedFile) {
                showToast('Please select an image first', 'error');
                return;
            }

            // Reset previous results
            outputPreview.style.display = 'none';
            outputContainer.querySelector('.placeholder').style.display = 'block';
            downloadSection.classList.remove('active');
            if (currentResultUrl) {
                URL.revokeObjectURL(currentResultUrl);
                currentResultUrl = null;
            }
            resetProgress();
            
            // Disable button during processing
            enhanceBtn.disabled = true;
            enhanceBtn.textContent = '⏳ Processing...';

            try {
                const formData = new FormData();
                formData.append('file', selectedFile);

                const response = await fetch('/convert', {
                    method: 'POST',
                    body: formData
                });

                if (!response.ok) {
                    throw new Error('Failed to start processing');
                }

                const data = await response.json();
                currentJobId = data.job_id;

                // Start polling for progress
                if (pollingInterval) {
                    clearInterval(pollingInterval);
                }

                pollingInterval = setInterval(async () => {
                    try {
                        const res = await fetch(`/progress/${currentJobId}`);
                        const progress = await res.json();

                        updateProgress(progress);

                        if (progress.progress >= 100) {
                            clearInterval(pollingInterval);
                            pollingInterval = null;
                            
                            // Fetch the result
                            const resultResponse = await fetch(`/result/${currentJobId}?t=${Date.now()}`);
                            if (resultResponse.ok) {
                                const blob = await resultResponse.blob();
                                if (currentResultUrl) {
                                    URL.revokeObjectURL(currentResultUrl);
                                }
                                currentResultUrl = URL.createObjectURL(blob);
                                outputPreview.src = currentResultUrl;
                                outputPreview.style.display = 'block';
                                outputContainer.querySelector('.placeholder').style.display = 'none';
                                
                                // Show download section
                                downloadSection.classList.add('active');
                                
                                showToast('✨ Enhancement completed!', 'success');
                            }
                            
                            enhanceBtn.disabled = false;
                            enhanceBtn.textContent = '🚀 Enhance Face';
                        } else if (progress.progress < 0) {
                            // Error occurred
                            clearInterval(pollingInterval);
                            pollingInterval = null;
                            showToast(`Error: ${progress.message}`, 'error');
                            enhanceBtn.disabled = false;
                            enhanceBtn.textContent = '🚀 Enhance Face';
                        }
                    } catch (error) {
                        console.error('Polling error:', error);
                    }
                }, 500);

            } catch (error) {
                console.error('Enhancement error:', error);
                showToast(`Error: ${error.message}`, 'error');
                enhanceBtn.disabled = false;
                enhanceBtn.textContent = '🚀 Enhance Face';
            }
        }

        // Reset everything
        function resetAll() {
            removeSelectedFile();
            if (pollingInterval) {
                clearInterval(pollingInterval);
                pollingInterval = null;
            }
            resetProgress();
            outputPreview.style.display = 'none';
            outputContainer.querySelector('.placeholder').style.display = 'block';
            downloadSection.classList.remove('active');
            if (currentResultUrl) {
                URL.revokeObjectURL(currentResultUrl);
                currentResultUrl = null;
            }
            enhanceBtn.disabled = true;
            enhanceBtn.textContent = '🚀 Enhance Face';
            showToast('Reset complete', 'info');
        }

        // Event Listeners
        fileInput.addEventListener('change', function(e) {
            if (this.files && this.files[0]) {
                handleFileSelect(this.files[0]);
            }
        });

        uploadArea.addEventListener('click', function() {
            fileInput.click();
        });

        uploadArea.addEventListener('dragover', function(e) {
            e.preventDefault();
            this.classList.add('dragover');
        });

        uploadArea.addEventListener('dragleave', function(e) {
            e.preventDefault();
            this.classList.remove('dragover');
        });

        uploadArea.addEventListener('drop', function(e) {
            e.preventDefault();
            this.classList.remove('dragover');
            if (e.dataTransfer.files && e.dataTransfer.files[0]) {
                handleFileSelect(e.dataTransfer.files[0]);
                // Update file input
                const dt = new DataTransfer();
                dt.items.add(e.dataTransfer.files[0]);
                fileInput.files = dt.files;
            }
        });

        removeFile.addEventListener('click', removeSelectedFile);

        enhanceBtn.addEventListener('click', enhance);

        resetBtn.addEventListener('click', resetAll);

        downloadBtn.addEventListener('click', downloadImage);

        // Keyboard shortcuts
        document.addEventListener('keydown', function(e) {
            if (e.key === 'Enter' && !enhanceBtn.disabled) {
                enhance();
            }
            if (e.key === 'Escape') {
                resetAll();
            }
            if ((e.ctrlKey || e.metaKey) && e.key === 's' && downloadSection.classList.contains('active')) {
                e.preventDefault();
                downloadImage();
            }
        });

        // Initial state
        console.log('✨ Face Enhancer ready!');
        console.log('💡 Shortcuts: Enter to enhance, Escape to reset, Ctrl+S to download');
    </script>

</body>
</html>
"""

@app.get("/", response_class=HTMLResponse)
async def home():
    return HTML_PAGE


# ====================================================
# API Endpoints
# ====================================================

@app.post("/convert")
async def convert(file: UploadFile = File(...)):
    contents = await file.read()
    
    np_img = np.frombuffer(
        contents,
        np.uint8
    )
    
    image = cv2.imdecode(
        np_img,
        cv2.IMREAD_COLOR
    )
    
    image = cv2.cvtColor(
        image,
        cv2.COLOR_BGR2RGB
    )
    
    job_id = str(uuid.uuid4())
    
    jobs[job_id] = {
        "progress": 0,
        "message": "Queued"
    }
    
    threading.Thread(
        target=process_job,
        args=(job_id, image),
        daemon=True
    ).start()
    
    return {
        "job_id": job_id
    }

@app.get("/progress/{job_id}")
async def progress(job_id: str):
    return jobs.get(
        job_id,
        {
            "progress": 0,
            "message": "Unknown job"
        }
    )

@app.get("/result/{job_id}")
async def result(job_id: str):
    if job_id not in results:
        return {
            "status": "processing"
        }
    
    return StreamingResponse(
        BytesIO(results[job_id]),
        media_type="image/png",
        headers={
            "Content-Disposition": f"attachment; filename=enhanced_{job_id[:8]}.png"
        }
    )

@app.get("/download/{job_id}")
async def download_result(job_id: str):
    """Direct download endpoint that forces browser to download as PNG"""
    if job_id not in results:
        return {
            "status": "processing",
            "message": "Result not ready yet"
        }
    
    return StreamingResponse(
        BytesIO(results[job_id]),
        media_type="image/png",
        headers={
            "Content-Disposition": f"attachment; filename=enhanced_{job_id[:8]}.png",
            "Cache-Control": "no-cache, no-store, must-revalidate",
            "Pragma": "no-cache",
            "Expires": "0"
        }
    )

# ====================================================
# Run
# ====================================================

if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)