Upload 2 files

app.py

@@ -0,0 +1,498 @@
+import subprocess
+import torch
+from PIL import Image
+import requests
+from io import BytesIO
+from transformers import AutoProcessor, AutoModelForCausalLM
+import os
+import threading
+import time
+import urllib.parse
+from fastapi import FastAPI, UploadFile, File, HTTPException, Form
+from fastapi.responses import JSONResponse
+
+app = FastAPI(
+    title="Florence-2 Image Captioning Server",
+    description="Auto-captions images from middleware server using Florence-2"
+)
+import threading
+import time
+import urllib.parse
+
+# Attempt to install flash-attn
+try:
+    subprocess.run('pip install flash-attn --no-build-isolation timm einops', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, check=True, shell=True)
+except subprocess.CalledProcessError as e:
+    print(f"Error installing flash-attn: {e}")
+    print("Continuing without flash-attn.")
+
+# Determine the device to use
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load the base model and processor
+try:
+    vision_language_model_base = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
+    vision_language_processor_base = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
+    print("✓ Base model loaded successfully")
+except Exception as e:
+    print(f"Error loading base model: {e}")
+    vision_language_model_base = None
+    vision_language_processor_base = None
+
+# Load the large model and processor
+try:
+    vision_language_model_large = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True).to(device).eval()
+    vision_language_processor_large = AutoProcessor.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True)
+    print("✓ Large model loaded successfully")
+except Exception as e:
+    print(f"Error loading large model: {e}")
+    vision_language_model_large = None
+    vision_language_processor_large = None
+
+def load_image_from_url(image_url):
+    """Load an image from a URL."""
+    try:
+        response = requests.get(image_url, timeout=30)
+        response.raise_for_status()
+        image = Image.open(BytesIO(response.content))
+        return image.convert('RGB')
+    except Exception as e:
+        raise ValueError(f"Error loading image from URL: {e}")
+
+def process_image_description(model, processor, image):
+    """Process an image and generate description using the specified model."""
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    inputs = processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=1024,
+            early_stopping=False,
+            do_sample=False,
+            num_beams=3,
+        )
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    processed_description = processor.post_process_generation(
+        generated_text,
+        task="<MORE_DETAILED_CAPTION>",
+        image_size=(image.width, image.height)
+    )
+    image_description = processed_description["<MORE_DETAILED_CAPTION>"]
+    return image_description
+
+def describe_image(uploaded_image, model_choice):
+    """Generate description from uploaded image."""
+    if uploaded_image is None:
+        return "Please upload an image."
+
+    if model_choice == "Florence-2-base":
+        if vision_language_model_base is None:
+            return "Base model failed to load."
+        model = vision_language_model_base
+        processor = vision_language_processor_base
+    elif model_choice == "Florence-2-large":
+        if vision_language_model_large is None:
+            return "Large model failed to load."
+        model = vision_language_model_large
+        processor = vision_language_processor_large
+    else:
+        return "Invalid model choice."
+
+    try:
+        return process_image_description(model, processor, uploaded_image)
+    except Exception as e:
+        return f"Error generating caption: {str(e)}"
+
+def describe_image_from_url(image_url, model_choice):
+    """Generate description from image URL."""
+    try:
+        if not image_url:
+            return {"error": "image_url is required"}
+        
+        if model_choice not in ["Florence-2-base", "Florence-2-large"]:
+            return {"error": "Invalid model choice. Use 'Florence-2-base' or 'Florence-2-large'"}
+        
+        # Load image from URL
+        image = load_image_from_url(image_url)
+        
+        # Select model and processor
+        if model_choice == "Florence-2-base":
+            if vision_language_model_base is None:
+                return {"error": "Base model not available"}
+            model = vision_language_model_base
+            processor = vision_language_processor_base
+        else:
+            if vision_language_model_large is None:
+                return {"error": "Large model not available"}
+            model = vision_language_model_large
+            processor = vision_language_processor_large
+        
+        # Generate caption
+        caption = process_image_description(model, processor, image)
+        
+        return {
+            "status": "success",
+            "model": model_choice,
+            "caption": caption,
+            "image_size": {"width": image.width, "height": image.height}
+        }
+        
+    except Exception as e:
+        return {"error": f"Error processing image: {str(e)}"}
+
+
+IMAGE_SERVER_BASE = os.getenv("IMAGE_SERVER_BASE", " ")
+DATA_COLLECTION_BASE = os.getenv("DATA_COLLECTION_BASE", "https://fred808-flow.hf.space")
+REQUESTER_ID = os.getenv("FLO_REQUESTER_ID", f"florence-2-{os.getpid()}")
+MODEL_CHOICE = os.getenv("FLO_MODEL_CHOICE", "Florence-2-base")
+
+
+def sanitize_name(name: str, max_len: int = 200) -> str:
+    """Sanitize a filename while preserving extension."""
+    import re
+    name = str(name).strip()
+    # replace spaces with underscores
+    name = re.sub(r"\s+", "_", name)
+    # remove any characters not alphanumeric, dot, dash, or underscore
+    name = re.sub(r"[^A-Za-z0-9_.-]", "", name)
+    if len(name) > max_len:
+        base, ext = os.path.splitext(name)
+        name = base[: max_len - len(ext)] + ext
+    return name or "file"
+
+def _build_download_url(course: str, video: str, frame: str) -> str:
+    """Build download URL with proper encoding of all path segments."""
+    # The middleware /download endpoint expects the 'file' parameter to be
+    # a path relative to the course folder (e.g. "video_name/frame.jpg").
+    # Frames live under a "{base_course}_frames" folder.
+    base_course = course
+    if not base_course.endswith("_frames"):
+        course_dir = f"{base_course}_frames"
+    else:
+        course_dir = base_course
+        base_course = course_dir[:-7]  # strip _frames for consistency
+
+    # Sanitize and encode path segments
+    safe_course = sanitize_name(course_dir)
+    safe_video = sanitize_name(video)
+    safe_frame = sanitize_name(frame)
+    
+    file_param = f"{safe_video}/{safe_frame}"
+    url = f"{IMAGE_SERVER_BASE.rstrip('/')}/download?course={urllib.parse.quote(safe_course, safe='')}&file={urllib.parse.quote(file_param, safe='')}"
+    print(f"[BACKGROUND] Built URL: {url}")
+    return url, safe_frame
+
+
+def _download_bytes(url: str, timeout: int = 30, chunk_size=32768):
+    try:
+        print(f"[BACKGROUND] Starting download: {url}")
+        response = requests.get(url, timeout=timeout, stream=True)
+        response.raise_for_status()
+        content = BytesIO()
+        total_size = int(response.headers.get('content-length', 0))
+        print(f"[BACKGROUND] Total size: {total_size} bytes")
+        
+        bytes_read = 0
+        for chunk in response.iter_content(chunk_size=chunk_size):
+            if chunk:
+                content.write(chunk)
+                bytes_read += len(chunk)
+                if total_size:
+                    print(f"\rDownloading: {bytes_read}/{total_size} bytes ({(bytes_read/total_size)*100:.1f}%)", end="", flush=True)
+        print()  # New line after progress
+        print(f"[BACKGROUND] Download complete: {bytes_read} bytes")
+        return content.getvalue(), response.headers.get('content-type')
+    except Exception as e:
+        print(f"[BACKGROUND] download failed {url}: {e}")
+        return None, None
+
+
+def _post_submit(caption: str, image_name: str, course: str, image_url: str, image_bytes: bytes):
+    submit_url = f"{DATA_COLLECTION_BASE.rstrip('/')}/submit"
+    files = {'image': (image_name, image_bytes, 'application/octet-stream')}
+    data = {'caption': caption, 'image_name': image_name, 'course': course, 'image_url': image_url}
+    
+    print(f"[BACKGROUND] Submitting to {submit_url}")
+    print(f"[BACKGROUND] Image name: {image_name}")
+    print(f"[BACKGROUND] Course: {course}")
+    print(f"[BACKGROUND] Caption length: {len(caption)} chars")
+    
+    try:
+        r = requests.post(submit_url, data=data, files=files, timeout=30)
+        print(f"[BACKGROUND] Submit response status: {r.status}")
+        try:
+            resp = r.json()
+            print(f"[BACKGROUND] Submit response JSON: {resp}")
+            return r.status_code, resp
+        except Exception:
+            print(f"[BACKGROUND] Submit response text: {r.text}")
+            return r.status_code, r.text
+    except Exception as e:
+        print(f"[BACKGROUND] Submit POST failed: {e}")
+        return None, None
+
+
+def _release_frame(course: str, video: str, frame: str):
+    try:
+        release_url = f"{IMAGE_SERVER_BASE.rstrip('/')}/middleware/release/frame/{urllib.parse.quote(course, safe='')}/{urllib.parse.quote(video, safe='')}/{urllib.parse.quote(frame, safe='')}"
+        requests.post(release_url, params={"requester_id": REQUESTER_ID}, timeout=10)
+    except Exception as e:
+        print(f"[BACKGROUND] release frame failed: {e}")
+
+
+def _release_course(course: str):
+    try:
+        release_url = f"{IMAGE_SERVER_BASE.rstrip('/')}/middleware/release/course/{urllib.parse.quote(course, safe='')}"
+        requests.post(release_url, params={"requester_id": REQUESTER_ID}, timeout=10)
+    except Exception as e:
+        print(f"[BACKGROUND] release course failed: {e}")
+
+
+# Background worker implementation
+def background_worker():
+    """Background worker that processes images from the middleware server."""
+    print("[BACKGROUND] Starting worker, waiting for model...")
+    
+    # Wait for model to be ready
+    waited = 0
+    while waited < 120:
+        if MODEL_CHOICE == "Florence-2-base" and vision_language_model_base:
+            break
+        elif MODEL_CHOICE == "Florence-2-large" and vision_language_model_large:
+            break
+        time.sleep(1)
+        waited += 1
+    
+    if waited >= 120:
+        print("[BACKGROUND] Model not available after timeout")
+        return
+
+    print(f"[BACKGROUND] Model {MODEL_CHOICE} ready, starting processing loop")
+    
+    while True:
+        try:
+            # Get next course
+            courses_url = f"{IMAGE_SERVER_BASE}/courses"
+            print(f"[BACKGROUND] Fetching courses from {courses_url}")
+            
+            try:
+                r = requests.get(courses_url, timeout=15)
+                r.raise_for_status()
+                courses_data = r.json()
+                
+                if not courses_data.get('courses'):
+                    print("[BACKGROUND] No courses found, waiting...")
+                    time.sleep(3)
+                    continue
+                
+                # Get first course
+                course_entry = courses_data['courses'][0]
+                if isinstance(course_entry, dict):
+                    course = course_entry.get('course_folder')
+                else:
+                    course = str(course_entry)
+                
+                if not course:
+                    print("[BACKGROUND] Invalid course entry")
+                    time.sleep(2)
+                    continue
+                
+                print(f"[BACKGROUND] Processing course: {course}")
+                
+                # Get images list
+                images_url = f"{IMAGE_SERVER_BASE}/images/{urllib.parse.quote(course, safe='')}"
+                r = requests.get(images_url, timeout=15)
+                r.raise_for_status()
+                images_data = r.json()
+                
+                if isinstance(images_data, dict):
+                    image_list = images_data.get('images', [])
+                else:
+                    image_list = images_data
+                
+                if not image_list:
+                    print(f"[BACKGROUND] No images found for course {course}")
+                    time.sleep(2)
+                    continue
+                
+                print(f"[BACKGROUND] Found {len(image_list)} images")
+                
+                # Process images
+                for img_entry in image_list:
+                    try:
+                        # Extract filename and metadata
+                        if isinstance(img_entry, dict):
+                            filename = img_entry.get('filename')
+                            if not filename:
+                                continue
+                        else:
+                            filename = str(img_entry)
+                        
+                        # Download image
+                        download_url = f"{IMAGE_SERVER_BASE}/images/{urllib.parse.quote(course, safe='')}/{urllib.parse.quote(filename, safe='')}"
+                        print(f"[BACKGROUND] Downloading {download_url}")
+                        
+                        img_bytes, _ = _download_bytes(download_url)
+                        if not img_bytes:
+                            print(f"[BACKGROUND] Failed to download {filename}")
+                            continue
+                        
+                        # Process with Florence
+                        try:
+                            pil_img = Image.open(BytesIO(img_bytes)).convert('RGB')
+                            
+                            if MODEL_CHOICE == "Florence-2-base":
+                                model = vision_language_model_base
+                                processor = vision_language_processor_base
+                            else:
+                                model = vision_language_model_large
+                                processor = vision_language_processor_large
+                            
+                            print(f"[BACKGROUND] Generating caption for {filename}")
+                            caption = process_image_description(model, processor, pil_img)
+                            print(f"[BACKGROUND] Generated caption for {filename}:")
+                            print("-" * 40)
+                            print(caption)
+                            print("-" * 40)
+                            
+                            # Submit result
+                            print(f"[BACKGROUND] Submitting caption to {DATA_COLLECTION_BASE}/submit")
+                            status, resp = _post_submit(caption, filename, course, download_url, img_bytes)
+                            if status and status < 400:
+                                print(f"[BACKGROUND] Successfully submitted {filename} (status={status})")
+                                if resp:
+                                    print(f"[BACKGROUND] Response: {resp}")
+                            else:
+                                print(f"[BACKGROUND] Failed to submit {filename}: status={status}, response={resp}")
+                            
+                        except Exception as e:
+                            print(f"[BACKGROUND] Error processing {filename}: {e}")
+                            continue
+                        finally:
+                            # Clean up
+                            if 'pil_img' in locals():
+                                del pil_img
+                            if 'img_bytes' in locals():
+                                del img_bytes
+                        
+                        time.sleep(0.5)  # Small delay between images
+                        
+                    except Exception as e:
+                        print(f"[BACKGROUND] Error in image loop: {e}")
+                        continue
+                
+                print(f"[BACKGROUND] Completed course {course}")
+                time.sleep(1)
+                
+            except Exception as e:
+                print(f"[BACKGROUND] Error in course loop: {e}")
+                time.sleep(5)
+                continue
+            
+        except Exception as e:
+            print(f"[BACKGROUND] Main loop error: {e}")
+            time.sleep(5)
+
+
+def _start_worker_thread():
+    """Start the background worker thread."""
+    t = threading.Thread(target=background_worker, daemon=True)
+    t.start()
+    return t
+
+
+# FastAPI endpoints for status/health
+@app.get("/")
+async def root():
+    return {
+        "name": "Florence-2 Image Captioning Server",
+        "status": "running",
+        "model_base": vision_language_model_base is not None,
+        "model_large": vision_language_model_large is not None,
+        "device": device
+    }
+
+@app.get("/health")
+async def health():
+    return {
+        "status": "healthy",
+        "model_base": vision_language_model_base is not None,
+        "model_large": vision_language_model_large is not None,
+        "device": device,
+        "model_choice": MODEL_CHOICE
+    }
+
+# Start background worker thread (daemon) so it doesn't block shutdown
+def _start_worker_thread():
+    t = threading.Thread(target=background_worker, daemon=True)
+    t.start()
+
+# Start background worker when FastAPI starts
+@app.on_event("startup")
+async def startup_event():
+    _start_worker_thread()
+
+
+@app.get("/analyze")
+async def analyze_get(image_url: str = None, model_choice: str = None):
+    """Analyze an image by URL. Usage: /analyze?image_url=https://...&model_choice=Florence-2-base"""
+    try:
+        mc = model_choice or MODEL_CHOICE
+        if image_url:
+            result = describe_image_from_url(image_url, mc)
+            if isinstance(result, dict) and result.get("status") == "success":
+                return JSONResponse(content={"success": True, "caption": result.get("caption"), "image_size": result.get("image_size")})
+            else:
+                return JSONResponse(status_code=400, content={"success": False, "error": result})
+        else:
+            raise HTTPException(status_code=400, detail="image_url query parameter is required")
+    except HTTPException:
+        raise
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})
+
+
+@app.post("/analyze")
+async def analyze_post(file: UploadFile = File(None), model_choice: str = Form(None)):
+    """Analyze an uploaded image (multipart/form-data). Returns caption JSON."""
+    try:
+        mc = model_choice or MODEL_CHOICE
+        if file is None:
+            raise HTTPException(status_code=400, detail="file is required")
+
+        content = await file.read()
+        try:
+            pil_img = Image.open(BytesIO(content)).convert('RGB')
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Failed to read uploaded image: {e}")
+
+        # Choose model
+        if mc == "Florence-2-large":
+            if vision_language_model_large is None:
+                raise HTTPException(status_code=503, detail="Base model not loaded")
+            model = vision_language_model_large
+            processor = vision_language_processor_large
+        else:
+            if vision_language_model_large is None:
+                raise HTTPException(status_code=503, detail="Large model not loaded")
+            model = vision_language_model_large
+            processor = vision_language_processor_large
+
+        caption = process_image_description(model, processor, pil_img)
+        return JSONResponse(content={"success": True, "caption": caption})
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"success": False, "error": str(e)})
+
+# Get the port from environment variable (for Hugging Face Spaces)
+port = int(os.environ.get("PORT", 7860))
+
+# Launch FastAPI with uvicorn when run directly
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=port)

requirements.txt

@@ -0,0 +1,5 @@
+transformers==4.48.0
+timm
+einops
+pillow
+hf_transfer