update app.py

app.py

@@ -21,12 +21,6 @@ st.set_page_config(
 HF_TOKEN = os.environ.get("HF_TOKEN", "")
 JINA_KEY = os.environ.get("JINA_KEY", "")
 
-QWEN_URL   = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-1.5B-Instruct/v1/chat/completions"
-HF_HEADERS = {
-    "Authorization": f"Bearer {HF_TOKEN}",
-    "Content-Type":  "application/json"
-}
-
 JINA_URL     = "https://api.jina.ai/v1/rerank"
 JINA_HEADERS = {
     "Authorization": f"Bearer {JINA_KEY}",
@@ -42,25 +36,30 @@ DETECT_PROMPT = (
     "jacket . dress . shirt . hat . bag ."
 )
 
-if not HF_TOKEN:
-    st.error("HF_TOKEN missing. Go to Space Settings → Secrets and add it.")
-    st.stop()
-
 if not JINA_KEY:
     st.error("JINA_KEY missing. Go to Space Settings → Secrets and add it.")
     st.stop()
 
+# ============================================================================
+# LOAD LOCAL MODELS
+# GIT-Large-COCO: caption generation
+# BLIP ITM:       image-text matching + cosine similarity
+# DINO:           object detection
+# Qwen2.5-1.5B:  caption fusion (moved local — API was returning 404)
+# ============================================================================
 @st.cache_resource
 def load_local_models():
     from transformers import (
         AutoProcessor,
         AutoModelForCausalLM,
+        AutoTokenizer,
         BlipProcessor,
         BlipForImageTextRetrieval,
         AutoModelForZeroShotObjectDetection
     )
     gc.collect()
 
+    # GIT-Large-COCO — caption generation
     git_processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
     git_model     = AutoModelForCausalLM.from_pretrained(
         "microsoft/git-large-coco",
@@ -68,6 +67,7 @@ def load_local_models():
     )
     git_model.eval()
 
+    # BLIP — ITM scoring and cosine similarity
     blip_processor = BlipProcessor.from_pretrained(
         "Salesforce/blip-image-captioning-large"
     )
@@ -77,6 +77,7 @@ def load_local_models():
     )
     blip_itm_model.eval()
 
+    # DINO — object detection
     dino_processor = AutoProcessor.from_pretrained(
         "IDEA-Research/grounding-dino-base"
     )
@@ -86,7 +87,22 @@ def load_local_models():
     )
     dino_model.eval()
 
-    return git_processor, git_model, blip_processor, blip_itm_model, dino_processor, dino_model
+    # Qwen2.5-1.5B — caption fusion (local, no API)
+    qwen_tokenizer = AutoTokenizer.from_pretrained(
+        "Qwen/Qwen2.5-1.5B-Instruct"
+    )
+    qwen_model = AutoModelForCausalLM.from_pretrained(
+        "Qwen/Qwen2.5-1.5B-Instruct",
+        torch_dtype=torch.float32
+    )
+    qwen_model.eval()
+
+    return (
+        git_processor, git_model,
+        blip_processor, blip_itm_model,
+        dino_processor, dino_model,
+        qwen_tokenizer, qwen_model
+    )
 
 def image_to_bytes(image: Image.Image) -> bytes:
     buf = BytesIO()
@@ -98,45 +114,14 @@ def image_to_data_uri(image: Image.Image) -> str:
     b64 = base64.b64encode(raw).decode()
     return f"data:image/jpeg;base64,{b64}"
 
-# ============================================================================
-# ONLY CHANGE: generate_captions_git
-# Fix: 5 different generation strategies instead of just max_new_tokens
-# Greedy / beam search / sampling with different temperatures
-# ============================================================================
 def generate_captions_git(image: Image.Image, git_proc, git_mod) -> list:
 
     strategies = [
-        # Greedy — short deterministic baseline
-        {
-            "max_new_tokens": 30
-        },
-        # Beam search — explores multiple decode paths
-        {
-            "max_new_tokens": 50,
-            "num_beams": 5,
-            "early_stopping": True
-        },
-        # Sampling — low temperature, focused output
-        {
-            "max_new_tokens": 60,
-            "do_sample": True,
-            "temperature": 0.7,
-            "top_k": 50
-        },
-        # Sampling — high temperature, creative output
-        {
-            "max_new_tokens": 70,
-            "do_sample": True,
-            "temperature": 1.3,
-            "top_k": 100
-        },
-        # Nucleus sampling — top-p based
-        {
-            "max_new_tokens": 55,
-            "do_sample": True,
-            "top_p": 0.9,
-            "temperature": 1.0
-        },
+        {"max_new_tokens": 30},
+        {"max_new_tokens": 50, "num_beams": 5, "early_stopping": True},
+        {"max_new_tokens": 60, "do_sample": True, "temperature": 0.7, "top_k": 50},
+        {"max_new_tokens": 70, "do_sample": True, "temperature": 1.3, "top_k": 100},
+        {"max_new_tokens": 55, "do_sample": True, "top_p": 0.9, "temperature": 1.0},
     ]
 
     captions     = []
@@ -149,26 +134,20 @@ def generate_captions_git(image: Image.Image, git_proc, git_mod) -> list:
                     pixel_values=pixel_values,
                     **strategy
                 )
-
             cap = git_proc.batch_decode(
-                generated_ids,
-                skip_special_tokens=True
+                generated_ids, skip_special_tokens=True
             )[0].strip().lower()
-
             captions.append(cap if cap else "a scene shown in the image")
-
         except Exception as e:
             st.warning(f"GIT error: {str(e)[:80]}")
             captions.append("a scene shown in the image")
 
-    # Deduplicate while keeping order
     seen, unique = set(), []
     for c in captions:
         if c not in seen:
             seen.add(c)
             unique.append(c)
 
-    # If model still returns all duplicates keep originals so voting has input
     if len(unique) < 2:
         unique = captions
 
@@ -199,7 +178,6 @@ def compute_itm_scores(image, captions, blip_proc, blip_itm) -> list:
 def compute_jina_scores(image: Image.Image, captions: list) -> list:
     img_data_uri = image_to_data_uri(image)
     scores       = []
-
     for cap in captions:
         try:
             payload = {
@@ -209,10 +187,8 @@ def compute_jina_scores(image: Image.Image, captions: list) -> list:
                 "top_n":     1
             }
             response = requests.post(
-                JINA_URL,
-                headers=JINA_HEADERS,
-                json=payload,
-                timeout=30
+                JINA_URL, headers=JINA_HEADERS,
+                json=payload, timeout=30
             )
             if response.status_code == 200:
                 result = response.json()
@@ -251,7 +227,6 @@ def compute_cosine_scores(image, captions, blip_proc, blip_itm) -> list:
 
         sims = cosine_similarity(img_feat, cap_feat)[0]
         return [round(float(s), 4) for s in sims]
-
     except Exception as e:
         st.warning(f"Cosine error: {str(e)[:60]}")
         return [0.0] * len(captions)
@@ -283,9 +258,7 @@ def detect_objects(image, dino_proc, dino_mod, threshold=0.3) -> tuple:
 
         target_sizes = torch.tensor([image.size[::-1]])
         results      = dino_proc.post_process_grounded_object_detection(
-            outputs,
-            inputs.input_ids,
-            target_sizes=target_sizes
+            outputs, inputs.input_ids, target_sizes=target_sizes
         )[0]
 
         scores = results["scores"]
@@ -310,51 +283,66 @@ def detect_objects(image, dino_proc, dino_mod, threshold=0.3) -> tuple:
         ]
         formatted = "Detected objects: [" + ", ".join(sorted_labels) + "]"
         return formatted, sorted_labels
-
     except Exception as e:
         st.warning(f"DINO error: {str(e)[:80]}")
         return "Object detection unavailable", []
 
-def fuse_captions(cap1: str, cap2: str, objects: str) -> str:
+# ============================================================================
+# STEP 7 — QWEN2.5-1.5B (LOCAL): CAPTION FUSION
+# Moved from API to local — API was consistently returning 404
+# Uses chat template for proper instruct format
+# Prompt asks Qwen to enrich and add detail using detected objects
+# ============================================================================
+def fuse_captions(cap1: str, cap2: str, objects: str, qwen_tok, qwen_mod) -> str:
+
     system_prompt = (
         "You are an expert image captioning assistant. "
-        "Write ONE natural, fluent, descriptive caption combining the best details. "
-        "Return ONLY the caption, no explanation or prefix."
+        "Write ONE natural, fluent, detailed and descriptive caption. "
+        "Combine the best details from both captions and incorporate the detected objects. "
+        "Return ONLY the final caption, no explanation or prefix."
     )
     user_prompt = (
         f"Caption A: {cap1}\n"
         f"Caption B: {cap2}\n"
         f"{objects}\n\n"
-        "Fused caption:"
+        "Write a detailed fused caption:"
     )
+
     try:
-        payload = {
-            "model": "Qwen/Qwen2.5-1.5B-Instruct",
-            "messages": [
-                {"role": "system", "content": system_prompt},
-                {"role": "user",   "content": user_prompt}
-            ],
-            "max_tokens": 100,
-            "temperature": 0.3,
-            "top_p": 0.9
-        }
-        response = requests.post(
-            QWEN_URL,
-            headers=HF_HEADERS,
-            json=payload,
-            timeout=40
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user",   "content": user_prompt}
+        ]
+
+        text = qwen_tok.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
         )
-        if response.status_code == 200:
-            fused = response.json()["choices"][0]["message"]["content"].strip()
-            for prefix in ["Fused caption:", "Caption:", "Result:"]:
-                if fused.lower().startswith(prefix.lower()):
-                    fused = fused[len(prefix):].strip()
-            return fused if fused else cap1
-        else:
-            st.warning(f"Qwen API error {response.status_code}")
-            return cap1
+
+        model_inputs = qwen_tok([text], return_tensors="pt")
+
+        with torch.no_grad():
+            generated_ids = qwen_mod.generate(
+                **model_inputs,
+                max_new_tokens=120,
+                temperature=0.3,
+                do_sample=True,
+                top_p=0.9
+            )
+
+        # Strip input tokens from output
+        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):]
+        fused = qwen_tok.decode(output_ids, skip_special_tokens=True).strip()
+
+        for prefix in ["Fused caption:", "Caption:", "Result:", "Answer:"]:
+            if fused.lower().startswith(prefix.lower()):
+                fused = fused[len(prefix):].strip()
+
+        return fused if fused else cap1
+
     except Exception as e:
-        st.warning(f"Qwen exception: {str(e)[:60]}")
+        st.warning(f"Qwen fusion error: {str(e)[:80]}")
         return cap1
 
 with st.sidebar:
@@ -380,12 +368,12 @@ Best 2 captions selected
 **6. Grounding DINO** (Local)
 Object detection
 
-**7. Qwen2.5-1.5B** (API)
+**7. Qwen2.5-1.5B** (Local)
 Caption fusion
     """)
     st.markdown("---")
-    st.markdown("**Local:** GIT-Large, BLIP ITM, DINO")
-    st.markdown("**API:** Jina, Qwen2.5")
+    st.markdown("**Local:** GIT-Large, BLIP ITM, DINO, Qwen2.5")
+    st.markdown("**API:** Jina")
 
 st.title("Image Caption Fusion System")
 st.markdown("Upload an image to generate a refined, grounded caption.")
@@ -407,8 +395,13 @@ if uploaded_file is not None:
     with col_run:
         if st.button("Generate Caption", type="primary", use_container_width=True):
 
-            with st.spinner("Loading local models (first run takes 2-3 min)..."):
-                git_proc, git_mod, blip_proc, blip_itm, dino_proc, dino_mod = load_local_models()
+            with st.spinner("Loading local models (first run takes 3-4 min)..."):
+                (
+                    git_proc, git_mod,
+                    blip_proc, blip_itm,
+                    dino_proc, dino_mod,
+                    qwen_tok, qwen_mod
+                ) = load_local_models()
 
             progress = st.progress(0)
             status   = st.empty()
@@ -463,7 +456,7 @@ if uploaded_file is not None:
             st.write(" | ".join(obj_list) if obj_list else obj_str)
 
             status.info("Step 7/7: Fusing captions with Qwen2.5-1.5B...")
-            final = fuse_captions(best_1, best_2, obj_str)
+            final = fuse_captions(best_1, best_2, obj_str, qwen_tok, qwen_mod)
             progress.progress(100)
             status.success("Pipeline complete!")