Update app.py

app.py

@@ -2,11 +2,12 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from selenium import webdriver
-from selenium.webdriver.common.by import By
+from selenium.webdriver.chrome.service import Service
 from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
-from selenium.common.exceptions import TimeoutException
+from selenium.common.exceptions import TimeoutException, WebDriverException
 import time
 import logging
 import os
@@ -37,17 +38,24 @@ class TranscriptResponse(BaseModel):
 
 def init_driver():
     options = Options()
-    options.add_argument("--headless")
-    options.add_argument("--disable-gpu")
+    options.add_argument("--headless=new")
     options.add_argument("--no-sandbox")
     options.add_argument("--disable-dev-shm-usage")
-    options.add_argument("--log-level=3")
-    options.add_experimental_option("excludeSwitches", ["enable-automation"])
-    options.add_experimental_option('useAutomationExtension', False)
+    options.add_argument("--disable-gpu")
+    options.add_argument("--window-size=1280,1696")
     
-    # For Hugging Face Spaces
+    # Chrome binary location in Hugging Face Spaces
     options.binary_location = "/usr/bin/google-chrome"
-    return webdriver.Chrome(options=options)
+    
+    # ChromeDriver path
+    service = Service(executable_path="/usr/bin/chromedriver")
+    
+    try:
+        driver = webdriver.Chrome(service=service, options=options)
+        return driver
+    except WebDriverException as e:
+        logger.error(f"Driver initialization failed: {str(e)}")
+        raise
 
 @app.post("/transcript", response_model=TranscriptResponse)
 async def get_transcript(request: VideoRequest):
@@ -71,6 +79,7 @@ async def get_transcript(request: VideoRequest):
             cookie_button.click()
             logger.info("Accepted cookies")
         except TimeoutException:
+            logger.info("No cookie consent found")
             pass
 
         # Click more button
@@ -92,8 +101,14 @@ async def get_transcript(request: VideoRequest):
 
         # Extract transcript
         segments = driver.find_elements(By.CSS_SELECTOR, "div.ytd-transcript-segment-renderer")
-        transcript = [segment.find_element(By.CLASS_NAME, "segment-text").text.strip() 
-                     for segment in segments if segment.find_element(By.CLASS_NAME, "segment-text").text.strip()]
+        transcript = []
+        for segment in segments:
+            try:
+                text = segment.find_element(By.CLASS_NAME, "segment-text").text.strip()
+                if text:
+                    transcript.append(text)
+            except:
+                continue
 
         if not transcript:
             raise HTTPException(status_code=404, detail="No transcript available")
@@ -106,11 +121,12 @@ async def get_transcript(request: VideoRequest):
         )
 
     except TimeoutException as e:
-        logger.error(f"Timeout: {str(e)}")
+        error_msg = "Timed out waiting for page elements - the video might not have transcripts"
+        logger.error(error_msg)
         return TranscriptResponse(
             success=False,
             transcript=None,
-            error="Timed out waiting for page elements",
+            error=error_msg,
             processing_time=time.time() - start_time
         )
     except Exception as e: