Update app.py

app.py

@@ -8,7 +8,7 @@ import numpy as np
 import torch
 
 from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.responses import HTMLResponse, StreamingResponse, FileResponse
 
 from huggingface_hub import hf_hub_download
 from torchvision.transforms.functional import normalize
@@ -370,6 +370,23 @@ HTML_PAGE = """
             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;
@@ -520,6 +537,10 @@ HTML_PAGE = """
             font-weight: 600;
         }
 
+        .image-label .badge-success {
+            background: #10b981;
+        }
+
         .image-container {
             width: 100%;
             aspect-ratio: 1;
@@ -642,6 +663,35 @@ HTML_PAGE = """
             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;
@@ -669,6 +719,11 @@ HTML_PAGE = """
                 right: 20px;
                 max-width: none;
             }
+
+            .download-section {
+                flex-direction: column;
+                align-items: stretch;
+            }
         }
     </style>
 </head>
@@ -716,6 +771,17 @@ HTML_PAGE = """
                 </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">
@@ -735,7 +801,7 @@ HTML_PAGE = """
                 <div class="image-card">
                     <div class="image-label">
                         <span>✨ Enhanced</span>
-                        <span class="badge">Output</span>
+                        <span class="badge badge-success">Output</span>
                     </div>
                     <div class="image-container" id="outputContainer">
                         <div class="placeholder">
@@ -762,6 +828,8 @@ HTML_PAGE = """
         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');
@@ -775,6 +843,7 @@ HTML_PAGE = """
         let selectedFile = null;
         let pollingInterval = null;
         let currentJobId = null;
+        let currentResultUrl = null;
 
         // Toast function
         function showToast(message, type = 'info') {
@@ -839,6 +908,11 @@ HTML_PAGE = """
             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);
@@ -868,6 +942,28 @@ HTML_PAGE = """
             }
         }
 
+        // 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) {
@@ -878,6 +974,11 @@ HTML_PAGE = """
             // 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
@@ -920,10 +1021,17 @@ HTML_PAGE = """
                             const resultResponse = await fetch(`/result/${currentJobId}?t=${Date.now()}`);
                             if (resultResponse.ok) {
                                 const blob = await resultResponse.blob();
-                                const url = URL.createObjectURL(blob);
-                                outputPreview.src = url;
+                                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');
                             }
                             
@@ -960,6 +1068,11 @@ HTML_PAGE = """
             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');
@@ -1004,6 +1117,8 @@ HTML_PAGE = """
 
         resetBtn.addEventListener('click', resetAll);
 
+        downloadBtn.addEventListener('click', downloadImage);
+
         // Keyboard shortcuts
         document.addEventListener('keydown', function(e) {
             if (e.key === 'Enter' && !enhanceBtn.disabled) {
@@ -1012,10 +1127,15 @@ HTML_PAGE = """
             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>
@@ -1086,7 +1206,30 @@ async def result(job_id: str):
     
     return StreamingResponse(
         BytesIO(results[job_id]),
-        media_type="image/png"
+        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"
+        }
     )
 
 # ====================================================