Update app.py

app.py

@@ -1,5 +1,7 @@
 import os
-import torch
+import sys
+import subprocess
+import importlib
 import requests
 from PIL import Image
 from io import BytesIO
@@ -8,12 +10,59 @@ from pydantic import BaseModel, HttpUrl
 from transformers import AutoProcessor, AutoModelForCausalLM
 import uvicorn
 
+# ===== RUNTIME DEPENDENCY ENSURER =====
+# Hardcoded torch version to ensure compatibility at startup.
+REQUIRED_TORCH_VERSION = os.getenv("REQUIRED_TORCH_VERSION", "2.2.2")
+
+def ensure_torch_installed(required_version: str = REQUIRED_TORCH_VERSION):
+    """Ensure the required torch version is installed at runtime.
+    This will attempt to import torch and compare versions. If missing or different,
+    it will pip-install the requested version using the running Python executable.
+
+    Note: Installing torch at every start may be slow and may require build artifacts
+    specific to the platform. This helper uses a simple pip install; if your target
+    platform requires a special wheel or extra index URL, set up the environment
+    outside of this script or modify the install command accordingly.
+    """
+    try:
+        import torch as _t
+        v = getattr(_t, "__version__", "")
+        # match major.minor.patch prefix
+        if v and v.startswith(required_version):
+            print(f"[INFO] torch {v} already installed")
+            return _t
+        else:
+            print(f"[INFO] torch version {v} != {required_version}, will reinstall")
+    except Exception:
+        print("[INFO] torch not found, installing now")
+
+    cmd = [sys.executable, "-m", "pip", "install", f"torch=={required_version}"]
+    print(f"[INFO] Running: {' '.join(cmd)}")
+    try:
+        subprocess.check_call(cmd)
+        importlib.invalidate_caches()
+        import torch as _t2
+        print(f"[INFO] Installed torch {_t2.__version__}")
+        return _t2
+    except subprocess.CalledProcessError as e:
+        print(f"[ERROR] pip install failed: {e}")
+        raise
+
+
+# Ensure torch is available before using the model
+torch = ensure_torch_installed()
+
 # ===== CONFIG =====
 DEVICE = "cpu"                      # Use CPU for compatibility
 RESIZE_DIM = (512, 512)             # Resize images to this resolution
 MAX_IMAGE_SIZE = 10 * 1024 * 1024   # 10MB max image size
 TASK = "<MORE_DETAILED_CAPTION>"    # Hardcoded task
 
+# URL template for frame iteration - replace with your actual URL
+BASE_URL_TEMPLATE = "https://example.com/frames/frame_{frame}.jpg"
+START_FRAME = 1          # Starting frame number
+FRAME_PADDING = 6       # Number of digits to pad frame numbers with
+
 # ===== FastAPI App =====
 app = FastAPI(
     title="Florence-2 Image Analysis API",
@@ -81,6 +130,48 @@ def download_image(url: str) -> Image.Image:
     except Exception as e:
         raise ValueError(f"Failed to process image: {e}")
 
+
+def iterate_and_analyze(base_url_template: str, start: int = 1, padding: int = 6):
+    """Iterate over a templated frame URL and analyze images sequentially.
+
+    base_url_template should contain a placeholder `{frame}` which will be replaced by
+    the zero-padded frame number, for example:
+      https://example.com/download?course=XYZ&file=frame%3AXYZ%2F{frame}%2Fframe_000{n}.jpg
+
+    The function yields tuples: (frame_number, url, caption_or_error)
+    Continues until a frame fails to download (e.g., 404 error)
+    """
+    if "{frame}" not in base_url_template:
+        raise ValueError("base_url_template must contain '{frame}' placeholder")
+
+    consecutive_errors = 0
+    MAX_CONSECUTIVE_ERRORS = 3  # Stop after this many consecutive errors
+    
+    i = start
+    while True:
+        frame_id = str(i).zfill(padding)
+        url = base_url_template.format(frame=frame_id)
+        try:
+            img = download_image(url)
+            caption = analyze_image(img)
+            consecutive_errors = 0  # Reset error counter on success
+            yield (i, url, {"success": True, "caption": caption})
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 404:
+                print(f"[INFO] No more frames found after frame {i-1}")
+                break
+            yield (i, url, {"success": False, "error": str(e)})
+            consecutive_errors += 1
+        except Exception as e:
+            yield (i, url, {"success": False, "error": str(e)})
+            consecutive_errors += 1
+        
+        if consecutive_errors >= MAX_CONSECUTIVE_ERRORS:
+            print(f"[INFO] Stopping after {MAX_CONSECUTIVE_ERRORS} consecutive errors")
+            break
+            
+        i += 1
+
 def analyze_image(image: Image.Image) -> str:
     """Analyze image using Florence-2 model with hardcoded task"""
     if not processor or not model:
@@ -217,15 +308,53 @@ async def analyze_image_get(image_url: str):
 
 # ===== Main Execution =====
 if __name__ == "__main__":
-    port = int(os.getenv("PORT", 7860))
-    print(f"[INFO] Starting server on port {port}")
-    print(f"[INFO] Model status: {'Loaded' if (processor and model) else 'Failed to load'}")
-    print(f"[INFO] Task: {TASK}")
-    print(f"[INFO] API Documentation: http://localhost:{port}/docs")
-    
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=port,
-        reload=False
-    )
+    if not processor or not model:
+        print("[ERROR] Model failed to load. Cannot proceed with frame analysis.")
+        sys.exit(1)
+
+    print("[INFO] Starting frame analysis...")
+    print(f"[INFO] Using URL template: {BASE_URL_TEMPLATE}")
+    print(f"[INFO] Starting from frame {START_FRAME} with {FRAME_PADDING} digit padding")
+
+    results = []
+    for frame_num, url, result in iterate_and_analyze(
+        BASE_URL_TEMPLATE,
+        start=START_FRAME,
+        padding=FRAME_PADDING
+    ):
+        if result["success"]:
+            print(f"[SUCCESS] Frame {frame_num}: {result['caption']}")
+            results.append({
+                "frame": frame_num,
+                "url": url,
+                "caption": result["caption"]
+            })
+        else:
+            print(f"[ERROR] Frame {frame_num}: {result['error']}")
+            results.append({
+                "frame": frame_num,
+                "url": url,
+                "error": result["error"]
+            })
+
+    # Save results to a JSON file
+    import json
+    output_file = "frame_analysis_results.json"
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(results, f, indent=2, ensure_ascii=False)
+    print(f"[INFO] Results saved to {output_file}")
+
+    # Optional: start the API server after frame analysis
+    start_server = os.getenv("START_SERVER", "false").lower() == "true"
+    if start_server:
+        port = int(os.getenv("PORT", 7860))
+        print(f"[INFO] Starting server on port {port}")
+        print(f"[INFO] Task: {TASK}")
+        print(f"[INFO] API Documentation: http://localhost:{port}/docs")
+        
+        uvicorn.run(
+            app,
+            host="0.0.0.0",
+            port=port,
+            reload=False
+        )