Update app.py

app.py

@@ -5,9 +5,9 @@ import uuid
 import json
 import base64
 import random
+import zipfile
 import threading
 import concurrent.futures
-import subprocess
 from pathlib import Path
 from typing import List, Optional
 
@@ -18,7 +18,7 @@ from PIL import Image
 
 from gradio import Server
 from fastapi import Request, UploadFile, File, Form
-from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
+from fastapi.responses import HTMLResponse, JSONResponse, FileResponse, StreamingResponse
 from diffusers import Flux2KleinPipeline, AutoencoderKLFlux2
 
 # --- App Configuration & Directories ---
@@ -214,6 +214,32 @@ async def download_file(filename: str):
         return JSONResponse({"error": "File not found"}, status_code=404)
     return FileResponse(path, filename=filename, media_type="image/png")
 
+@app.get("/api/download-zip")
+async def download_zip(std: str, small: str):
+    """Packages both generated images into a single ZIP file and streams it."""
+    # Prevent path traversal by taking only the filename
+    std_name = Path(std).name
+    small_name = Path(small).name
+    
+    std_path = OUTPUT_DIR / std_name
+    small_path = OUTPUT_DIR / small_name
+    
+    if not std_path.exists() or not small_path.exists():
+        return JSONResponse({"error": "Generated files not found"}, status_code=404)
+        
+    memory_file = io.BytesIO()
+    with zipfile.ZipFile(memory_file, 'w', zipfile.ZIP_DEFLATED) as zf:
+        zf.write(std_path, arcname=f"Standard_Decoder_{std_name}")
+        zf.write(small_path, arcname=f"Small_Decoder_{small_name}")
+        
+    memory_file.seek(0)
+    
+    return StreamingResponse(
+        memory_file,
+        media_type="application/zip",
+        headers={"Content-Disposition": f"attachment; filename=Flux2_Comparison_{uuid.uuid4().hex[:6]}.zip"}
+    )
+
 @app.post("/api/compare")
 async def compare_images(
     prompt: str = Form(...),
@@ -258,6 +284,8 @@ async def compare_images(
             "seed": used_seed,
             "std_url": f"/download/{std_filename}",
             "small_url": f"/download/{small_filename}",
+            "std_filename": std_filename,
+            "small_filename": small_filename,
             "device": DEVICE_LABEL,
         })
 
@@ -365,6 +393,9 @@ async def homepage(request: Request):
       font-weight: 500;
       font-size: 1.1rem;
       flex-shrink: 0;
+      display: flex;
+      justify-content: space-between;
+      align-items: center;
     }}
 
     .panel-body-scroll {{
@@ -437,7 +468,6 @@ async def homepage(request: Request):
       cursor: pointer; font-size: 12px;
     }}
     
-    /* The (+) Button */
     .add-more-btn {{
       display: flex; align-items: center; justify-content: center;
       border: 2px dashed var(--ub-muted); border-radius: 4px;
@@ -461,40 +491,23 @@ async def homepage(request: Request):
     .advanced-body.open {{ display: block; }}
     .grid-2 {{ display: grid; grid-template-columns: 1fr 1fr; gap: 15px; }}
 
-    /* Status Container (Execution Log) */
+    /* Status Container */
     .status-container {{
-      margin-top: 20px;
-      margin-bottom: 20px;
-      border: 1px solid var(--ub-border);
-      border-radius: 4px;
-      background: #200014; /* deep ubuntu terminal background */
-      display: flex;
-      flex-direction: column;
-      flex: 1; 
-      min-height: 100px;
-      max-height: 200px; /* limits expansion if advanced is closed */
+      margin-top: 20px; margin-bottom: 20px;
+      border: 1px solid var(--ub-border); border-radius: 4px;
+      background: #200014; display: flex; flex-direction: column;
+      flex: 1; min-height: 100px; max-height: 200px;
     }}
     .status-header {{
-      padding: 8px 12px;
-      font-size: 11px;
-      font-weight: 700;
-      color: var(--ub-muted);
-      background: rgba(0,0,0,0.4);
-      border-bottom: 1px solid var(--ub-border);
-      text-transform: uppercase;
-      letter-spacing: 0.5px;
-      flex-shrink: 0;
+      padding: 8px 12px; font-size: 11px; font-weight: 700;
+      color: var(--ub-muted); background: rgba(0,0,0,0.4);
+      border-bottom: 1px solid var(--ub-border); text-transform: uppercase;
+      letter-spacing: 0.5px; flex-shrink: 0;
     }}
     .status-log {{
-      padding: 10px;
-      font-family: 'Courier New', Courier, monospace;
-      font-size: 12px;
-      color: #eeeeee;
-      overflow-y: auto;
-      flex: 1;
-      display: flex;
-      flex-direction: column;
-      gap: 4px;
+      padding: 10px; font-family: 'Courier New', Courier, monospace;
+      font-size: 12px; color: #eeeeee; overflow-y: auto;
+      flex: 1; display: flex; flex-direction: column; gap: 4px;
     }}
     .log-time {{ color: #777; margin-right: 8px; }}
     .log-info {{ color: #5bc0eb; }}
@@ -505,8 +518,7 @@ async def homepage(request: Request):
     .btn {{
       width: 100%; padding: 14px; border: none; border-radius: 4px;
       font-size: 16px; font-weight: 700; cursor: pointer;
-      transition: opacity 0.2s, background 0.2s;
-      flex-shrink: 0;
+      transition: opacity 0.2s, background 0.2s; flex-shrink: 0;
     }}
     .btn-primary {{
       background: var(--ub-orange); color: white;
@@ -515,42 +527,39 @@ async def homepage(request: Request):
     .btn-primary:hover {{ background: var(--ub-orange-hover); }}
     .btn:disabled {{ opacity: 0.6; cursor: not-allowed; }}
 
+    /* Top-Right Download Icon */
+    .action-icon {{
+      display: none; background: none; border: none; color: var(--ub-muted);
+      cursor: pointer; padding: 4px; transition: color 0.2s;
+    }}
+    .action-icon:hover {{ color: var(--ub-orange); }}
+
     /* SLIDER CONTAINER */
     .panel-body-slider {{
-      flex: 1;
-      display: flex;
-      flex-direction: column;
-      padding: 0; 
-      position: relative;
+      flex: 1; display: flex; flex-direction: column;
+      padding: 0; position: relative;
     }}
     .slider-stage {{
-      position: absolute;
-      top: 0; left: 0; right: 0; bottom: 0;
-      background: #111;
-      overflow: hidden;
-      display: flex;
-      align-items: center;
-      justify-content: center;
+      position: absolute; top: 0; left: 0; right: 0; bottom: 0;
+      background: #111; overflow: hidden; display: flex;
+      align-items: center; justify-content: center;
     }}
     .slider-empty {{ color: var(--ub-muted); text-align: center; z-index: 1; }}
     
     .slider-img {{
-      position: absolute; top: 0; left: 0;
-      width: 100%; height: 100%; object-fit: contain;
-      display: none; user-select: none; -webkit-user-drag: none;
+      position: absolute; top: 0; left: 0; width: 100%; height: 100%;
+      object-fit: contain; display: none; user-select: none; -webkit-user-drag: none;
     }}
     #imgSmall {{ clip-path: inset(0 50% 0 0); }}
 
     .slider-handle {{
       position: absolute; left: 50%; top: 0; bottom: 0;
-      width: 4px; background: var(--ub-orange);
-      cursor: ew-resize; display: none; z-index: 10;
+      width: 4px; background: var(--ub-orange); cursor: ew-resize; display: none; z-index: 10;
     }}
     .slider-handle::after {{
       content: '◀ ▶'; position: absolute; top: 50%; left: 50%;
-      transform: translate(-50%, -50%);
-      width: 40px; height: 30px; background: var(--ub-orange);
-      color: white; border-radius: 15px;
+      transform: translate(-50%, -50%); width: 40px; height: 30px;
+      background: var(--ub-orange); color: white; border-radius: 15px;
       display: flex; align-items: center; justify-content: center;
       font-size: 10px; font-weight: bold; box-shadow: 0 2px 6px rgba(0,0,0,0.5);
     }}
@@ -567,8 +576,7 @@ async def homepage(request: Request):
 
     .loader {{
       position: absolute; inset: 0; background: rgba(0,0,0,0.7);
-      display: none; flex-direction: column;
-      align-items: center; justify-content: center; z-index: 20;
+      display: none; flex-direction: column; align-items: center; justify-content: center; z-index: 20;
     }}
     .spinner {{
       width: 40px; height: 40px; border: 4px solid rgba(255,255,255,0.2);
@@ -587,7 +595,6 @@ async def homepage(request: Request):
       cursor: pointer; transition: transform 0.2s, box-shadow 0.2s;
     }}
     .ex-card:hover {{ transform: translateY(-3px); box-shadow: 0 6px 16px rgba(0,0,0,0.3); }}
-    
     .ex-card-img-wrap {{ width: 100%; aspect-ratio: 1; display: flex; background: #000; }}
     .ex-card-img-wrap img {{ height: 100%; object-fit: cover; }}
     .ex-card p {{ padding: 12px; margin: 0; font-size: 13px; color: var(--ub-muted); line-height: 1.4; }}
@@ -675,7 +682,16 @@ async def homepage(request: Request):
       </div>
 
       <div class="panel">
-        <div class="panel-header">Comparison View</div>
+        <div class="panel-header">
+          <span>Comparison View</span>
+          <button id="downloadZipBtn" class="action-icon" title="Download Both Images (ZIP)">
+            <svg width="22" height="22" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" viewBox="0 0 24 24">
+              <path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"></path>
+              <polyline points="7 10 12 15 17 10"></polyline>
+              <line x1="12" y1="15" x2="12" y2="3"></line>
+            </svg>
+          </button>
+        </div>
         <div class="panel-body-slider">
           <div class="slider-stage" id="sliderStage">
             <div class="slider-empty" id="sliderEmpty">
@@ -713,6 +729,8 @@ async def homepage(request: Request):
   <script>
     const examples = {examples_json};
     let filesState = [];
+    let currentStdFilename = "";
+    let currentSmallFilename = "";
 
     // UI Elements
     const dropZone = document.getElementById('dropZone');
@@ -721,6 +739,7 @@ async def homepage(request: Request):
     const uploadText = document.getElementById('uploadText');
     const promptInput = document.getElementById('promptInput');
     const runBtn = document.getElementById('runBtn');
+    const downloadZipBtn = document.getElementById('downloadZipBtn');
     
     // Status Log
     const statusLog = document.getElementById('statusLog');
@@ -807,10 +826,8 @@ async def homepage(request: Request):
         }}
         renderPreviews();
         
-        // Scroll to top
         window.scrollTo({{top: 0, behavior: 'smooth'}});
         
-        // Wait a tiny bit for UI update, then trigger generation
         setTimeout(() => {{
            logMsg("Example loaded. Starting comparison...", "log-info");
            runBtn.click();
@@ -829,7 +846,6 @@ async def homepage(request: Request):
       
       let imgHTML = '';
       if(ex.urls.length > 1) {{
-        // Split view for multiple image examples
         imgHTML = `
           <div class="ex-card-img-wrap">
              <img src="${{ex.urls[0]}}" style="width:50%; border-right:1px solid #000;">
@@ -871,6 +887,13 @@ async def homepage(request: Request):
       updateSlider(e.touches[0].clientX);
     }});
 
+    // --- Download Zip Logic ---
+    downloadZipBtn.onclick = () => {{
+      if(!currentStdFilename || !currentSmallFilename) return;
+      logMsg("Initiating ZIP download...", "log-info");
+      window.location.href = `/api/download-zip?std=${{currentStdFilename}}&small=${{currentSmallFilename}}`;
+    }};
+
     // --- Form Submission ---
     runBtn.onclick = async () => {{
       const prompt = promptInput.value.trim();
@@ -894,6 +917,7 @@ async def homepage(request: Request):
 
       loader.style.display = 'flex';
       runBtn.disabled = true;
+      downloadZipBtn.style.display = 'none';
 
       logMsg("Sending request to backend. Running both VAE models...", "log-info");
 
@@ -904,6 +928,9 @@ async def homepage(request: Request):
         if(data.success) {{
           logMsg(`Success! Inference completed. Used seed: ${{data.seed}}`, "log-success");
           
+          currentStdFilename = data.std_filename;
+          currentSmallFilename = data.small_filename;
+
           imgStd.src = data.std_url;
           imgSmall.src = data.small_url;
           
@@ -913,6 +940,7 @@ async def homepage(request: Request):
             imgSmall.style.display = 'block';
             sliderHandle.style.display = 'block';
             sliderLabels.style.display = 'flex';
+            downloadZipBtn.style.display = 'block'; // Reveal download button
             
             // Reset slider to center
             const rect = sliderStage.getBoundingClientRect();