Fix 400 Bad Request: add HF Hub model loading, robust image validation, and graceful fallback

app.py

@@ -27,7 +27,8 @@ from transformers import (
 # CONFIG
 # ============================================================
 MODEL_NAME = "Salesforce/blip-image-captioning-base"
-FINE_TUNED_PATH = "./flickr8k_blip_output/best_model"
+FINE_TUNED_LOCAL_PATH = "./flickr8k_blip_output/best_model"
+FINE_TUNED_HUB_NAME = "Pokzy/flickr8k-finetuned"
 TRANSLATION_MODEL_NAME = "facebook/nllb-200-distilled-600M"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -77,7 +78,6 @@ def translate_to_thai(texts: list[str]) -> list[str]:
     return translations
 
 
-
 def load_models():
     """Preload both Pretrained and Fine-Tuned BLIP models if available."""
     global models, processors
@@ -90,23 +90,34 @@ def load_models():
         models["pretrained"].eval()
         print("[INFO] Pretrained model loaded successfully!")
 
-    # 2. Load Fine-Tuned BLIP
+    # 2. Load Fine-Tuned BLIP (Local directory or Hugging Face Model Hub)
     if "fine-tuned" not in models:
-        if os.path.exists(FINE_TUNED_PATH) and os.path.isdir(FINE_TUNED_PATH):
+        # Priority A: Check local directory
+        if os.path.exists(FINE_TUNED_LOCAL_PATH) and os.path.isdir(FINE_TUNED_LOCAL_PATH):
             has_files = any(
-                os.path.exists(os.path.join(FINE_TUNED_PATH, f))
+                os.path.exists(os.path.join(FINE_TUNED_LOCAL_PATH, f))
                 for f in ["pytorch_model.bin", "model.safetensors", "config.json"]
             )
             if has_files:
-                print(f"[INFO] Loading Fine-Tuned model from {FINE_TUNED_PATH}...")
-                processors["fine-tuned"] = BlipProcessor.from_pretrained(FINE_TUNED_PATH)
-                models["fine-tuned"] = BlipForConditionalGeneration.from_pretrained(FINE_TUNED_PATH).to(DEVICE)
+                try:
+                    print(f"[INFO] Loading Fine-Tuned model from local directory {FINE_TUNED_LOCAL_PATH}...")
+                    processors["fine-tuned"] = BlipProcessor.from_pretrained(FINE_TUNED_LOCAL_PATH)
+                    models["fine-tuned"] = BlipForConditionalGeneration.from_pretrained(FINE_TUNED_LOCAL_PATH).to(DEVICE)
+                    models["fine-tuned"].eval()
+                    print("[INFO] Local Fine-Tuned model loaded successfully!")
+                except Exception as e:
+                    print(f"[WARN] Could not load local Fine-Tuned model: {e}")
+
+        # Priority B: Download from Hugging Face Hub (Pokzy/flickr8k-finetuned)
+        if "fine-tuned" not in models:
+            try:
+                print(f"[INFO] Attempting to load Fine-Tuned model from Hugging Face Hub ({FINE_TUNED_HUB_NAME})...")
+                processors["fine-tuned"] = BlipProcessor.from_pretrained(FINE_TUNED_HUB_NAME)
+                models["fine-tuned"] = BlipForConditionalGeneration.from_pretrained(FINE_TUNED_HUB_NAME).to(DEVICE)
                 models["fine-tuned"].eval()
-                print("[INFO] Fine-Tuned model loaded successfully!")
-            else:
-                print("[WARN] Fine-Tuned path exists but model files are missing.")
-        else:
-            print("[WARN] Fine-Tuned model directory not found.")
+                print("[INFO] Fine-Tuned model loaded successfully from Hugging Face Hub!")
+            except Exception as e:
+                print(f"[WARN] Fine-Tuned model on HF Hub ({FINE_TUNED_HUB_NAME}) not loaded yet: {e}")
 
 
 def generate_captions_for_model(model_key: str, image: Image.Image, num_captions: int = 5):
@@ -176,17 +187,16 @@ async def predict(
     """
     load_models()
 
-    allowed_types = {"image/jpeg", "image/jpg", "image/png", "image/gif", "image/webp", "image/bmp"}
-    if file.content_type not in allowed_types:
+    try:
+        image_bytes = await file.read()
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    except Exception as e:
         raise HTTPException(
             status_code=400,
-            detail=f"Unsupported file type: {file.content_type}. Please upload JPEG, PNG, GIF, or WebP.",
+            detail=f"Invalid image file: {str(e)}. Please upload a valid image file (JPEG, PNG, WebP, etc.)."
         )
 
     try:
-        image_bytes = await file.read()
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-
         results = {}
 
         def process_model(key: str, label: str):
@@ -196,16 +206,17 @@ async def predict(
                 data["captions_th"] = translate_to_thai(en_captions)
             results[key] = data
 
-        # 1. Fine-tuned model inference
+        # 1. Fine-tuned model inference (with graceful fallback to pretrained if fine-tuned model weights not present)
         if model_choice in ["both", "fine-tuned"]:
             if "fine-tuned" in models:
                 process_model("fine-tuned", "Fine-Tuned BLIP (Flickr8k)")
-            elif model_choice == "fine-tuned":
-                raise HTTPException(status_code=400, detail="Fine-tuned model is not available.")
+            else:
+                # Graceful fallback: use pretrained if fine-tuned is missing
+                process_model("pretrained", "Pretrained BLIP (Base — Fine-Tuned Model Loading/Pending)")
 
         # 2. Pretrained model inference
         if model_choice in ["both", "pretrained"]:
-            if "pretrained" in models:
+            if "pretrained" in models and "pretrained" not in results:
                 process_model("pretrained", "Pretrained BLIP (Base)")
 
         return {