Update app.py

app.py

@@ -10,7 +10,6 @@ from concurrent.futures import ThreadPoolExecutor
 
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 import rarfile
@@ -23,7 +22,6 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# 1. Enable CORS (good practice, though less critical if serving static from same origin)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -32,7 +30,8 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# 2. Utility Functions
+# --- Utility Functions ---
+
 def parse_srt_time_to_ms(time_str):
     try:
         if not time_str: return 0
@@ -80,16 +79,17 @@ def compress_image(image_bytes, quality=70, max_width=800):
             img = img.resize((width, height), Image.Resampling.LANCZOS)
         
         buffer = io.BytesIO()
-        img.save(buffer, format="JPEG", quality=quality)
+        # Quality expects integer 1-95
+        img.save(buffer, format="JPEG", quality=int(quality))
         return buffer.getvalue()
     except Exception as e:
         logger.error(f"Image compression failed: {e}")
         return None
 
-def process_batch_gemini(api_key, items):
+def process_batch_gemini(api_key, items, model_name):
     try:
         genai.configure(api_key=api_key)
-        model = genai.GenerativeModel('gemini-3-flash-preview') 
+        model = genai.GenerativeModel(model_name) 
 
         prompt_parts = [
             "You are a Subtitle Quality Control (QC) bot.",
@@ -114,31 +114,35 @@ def process_batch_gemini(api_key, items):
         logger.error(f"Gemini API Error with key ...{api_key[-4:]}: {e}")
         return None 
 
-# 3. API Endpoint
+# --- Main Endpoint ---
+
 @app.post("/api/analyze")
 async def analyze_subtitles(
     srt_file: UploadFile = File(...),
     media_files: list[UploadFile] = File(...),
     api_keys: str = Form(...),
-    batch_size: int = Form(20)
+    batch_size: int = Form(20),
+    model_name: str = Form("gemini-2.0-flash"),
+    compression_quality: float = Form(0.7) 
 ):
     temp_dir = tempfile.mkdtemp()
     try:
-        # Read SRT
+        # Convert float quality (0.1-1.0) to integer (10-100) for PIL
+        pil_quality = max(10, min(100, int(compression_quality * 100)))
+
+        # 1. Read SRT
         srt_content = (await srt_file.read()).decode('utf-8', errors='ignore')
         srt_data = parse_srt(srt_content)
         srt_data.sort(key=lambda x: x['startTimeMs'])
 
-        # Process Media
+        # 2. Process Media
         images = []
-        
         for file in media_files:
             file_path = os.path.join(temp_dir, file.filename)
             with open(file_path, "wb") as f:
                 shutil.copyfileobj(file.file, f)
             
             if file.filename.lower().endswith('.rar'):
-                # Docker guarantees 'unrar' is installed
                 try:
                     with rarfile.RarFile(file_path) as rf:
                         rf.extractall(temp_dir)
@@ -156,7 +160,7 @@ async def analyze_subtitles(
                     if ms is not None:
                         with open(full_path, "rb") as f:
                             raw_bytes = f.read()
-                            compressed = compress_image(raw_bytes)
+                            compressed = compress_image(raw_bytes, quality=pil_quality)
                             if compressed:
                                 images.append({
                                     "filename": filename,
@@ -166,12 +170,14 @@ async def analyze_subtitles(
 
         images.sort(key=lambda x: x['timeMs'])
         
+        # 3. Pair
         pairs = []
         for i in range(len(images)):
             img = images[i]
             srt = srt_data[i] if i < len(srt_data) else None
             
             if srt:
+                # Create Thumbnail (lower quality for UI speed)
                 thumb_bytes = compress_image(img['data'], quality=50, max_width=300)
                 thumb_b64 = base64.b64encode(thumb_bytes).decode('utf-8')
                 
@@ -189,7 +195,7 @@ async def analyze_subtitles(
         if not pairs:
             return {"status": "error", "message": "No valid image/subtitle pairs found."}
 
-        # Process Gemini
+        # 4. Process Gemini
         keys = [k.strip() for k in api_keys.split('\n') if k.strip()]
         if not keys:
             raise HTTPException(status_code=400, detail="No API Keys provided")
@@ -199,7 +205,7 @@ async def analyze_subtitles(
         
         def worker(batch_idx, batch):
             key = keys[batch_idx % len(keys)]
-            return process_batch_gemini(key, batch)
+            return process_batch_gemini(key, batch, model_name)
 
         with ThreadPoolExecutor(max_workers=len(keys)) as executor:
             futures = [executor.submit(worker, i, b) for i, b in enumerate(batches)]
@@ -209,6 +215,7 @@ async def analyze_subtitles(
                     for item in res:
                         results_map[item['index']] = item
 
+        # 5. Build Output
         final_output = []
         for p in pairs:
             analysis = results_map.get(p['index'])
@@ -240,5 +247,4 @@ async def analyze_subtitles(
     finally:
         shutil.rmtree(temp_dir)
 
-# 4. Serve Static Files (Frontend)
 app.mount("/", StaticFiles(directory="static", html=True), name="static")