Update app.py

app.py

@@ -1,18 +1,23 @@
 import os
 import json
-import requests
 from PIL import Image
 import torch
 import gradio as gr
+from transformers import (
+    BlipImageProcessor,
+    AutoTokenizer,
+    Llama4ForConditionalGeneration,
+)
 from ppt_parser import transfer_to_structure
-from transformers import AutoProcessor, Llama4ForConditionalGeneration
 
-# ✅ Hugging Face token
+# ✅ Load Hugging Face token
 hf_token = os.getenv("HF_TOKEN")
 model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
 
-# ✅ Load model & processor
-processor = AutoProcessor.from_pretrained(model_id, token=hf_token)
+# ✅ Load image processor, tokenizer, and model manually
+image_processor = BlipImageProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+
 model = Llama4ForConditionalGeneration.from_pretrained(
     model_id,
     token=hf_token,
@@ -21,7 +26,7 @@ model = Llama4ForConditionalGeneration.from_pretrained(
     torch_dtype=torch.bfloat16,
 )
 
-# ✅ Global storage
+# ✅ Global state
 extracted_text = ""
 image_paths = []
 
@@ -41,7 +46,6 @@ def extract_text_from_pptx_json(parsed_json: dict) -> str:
                         text += para.get("text", "") + "\n"
     return text.strip()
 
-# ✅ Handle uploaded PPTX
 def handle_pptx_upload(pptx_file):
     global extracted_text, image_paths
     tmp_path = pptx_file.name
@@ -50,57 +54,42 @@ def handle_pptx_upload(pptx_file):
     extracted_text = extract_text_from_pptx_json(parsed_json)
     return extracted_text or "No readable text found in slides."
 
-# ✅ Multimodal Q&A using Scout
 def ask_llama(question):
     global extracted_text, image_paths
 
     if not extracted_text and not image_paths:
-        return "Please upload and extract a PPTX first."
-
-    # 🧠 Build multimodal chat messages
-    messages = [
-        {
-            "role": "user",
-            "content": [],
-        }
-    ]
-
-    # Add up to 2 images to prevent OOM
-    for path in image_paths[:2]:
-        messages[0]["content"].append({"type": "image", "image": Image.open(path)})
-
-    messages[0]["content"].append({
-        "type": "text",
-        "text": f"{extracted_text}\n\nQuestion: {question}"
-    })
-
-    inputs = processor.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt"
-    ).to(model.device)
-
-    outputs = model.generate(**inputs, max_new_tokens=256)
-
-    response = processor.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:])[0]
+        return "Please upload and extract a PPTX file first."
+
+    # ✅ Use the first image only (you can expand to multiple with batching)
+    image = Image.open(image_paths[0]).convert("RGB")
+    vision_inputs = image_processor(images=image, return_tensors="pt").to(model.device)
+
+    prompt = f"<|user|>\n{extracted_text}\n\nQuestion: {question}<|end|>\n<|assistant|>\n"
+    text_inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output = model.generate(
+            input_ids=text_inputs["input_ids"],
+            pixel_values=vision_inputs["pixel_values"],
+            max_new_tokens=256,
+        )
+
+    response = tokenizer.decode(output[0][text_inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
     return response.strip()
 
 # ✅ Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Multimodal Llama 4 Scout Study Assistant")
+    gr.Markdown("## 🧠 Llama-4-Scout Multimodal Study Assistant")
 
     pptx_input = gr.File(label="📂 Upload PPTX File", file_types=[".pptx"])
-    extract_btn = gr.Button("📜 Extract Text + Images")
+    extract_btn = gr.Button("📜 Extract Text + Slides")
 
     extracted_output = gr.Textbox(label="📄 Slide Text", lines=10, interactive=False)
-
     extract_btn.click(handle_pptx_upload, inputs=[pptx_input], outputs=[extracted_output])
 
     question = gr.Textbox(label="❓ Ask a Question")
-    ask_btn = gr.Button("💬 Ask Llama 4 Scout")
-    ai_answer = gr.Textbox(label="🤖 Answer", lines=6)
+    ask_btn = gr.Button("💬 Ask Scout")
+    ai_answer = gr.Textbox(label="🤖 Llama Answer", lines=6)
 
     ask_btn.click(ask_llama, inputs=[question], outputs=[ai_answer])