# app.py — Robust CPU-friendly SigLip -> (Llava local | trust_remote_code | HF router) pipeline import os # Force CPU before importing torch/transformers if you want CPU-only os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") os.environ.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1") import sys import traceback import json from typing import List, Optional import requests import torch import torch.nn.functional as F from datasets import load_dataset from transformers import ( AutoProcessor, AutoModel, AutoTokenizer, AutoModelForCausalLM, ) from PIL import Image import gradio as gr from tqdm import tqdm # ------------------------- # Config - update these IDs as needed # ------------------------- SIGLIP_MODEL_ID = "EYEDOL/siglipFULL-agri-finetuned" LLAVA_MODEL_ID = "liuhaotian/llava-v1.6-vicuna-7b" DATASET_TEMPLATE = "EYEDOL/AGRILLAVA-image-text{}" NUM_DATASETS = 1 # set to 15 if you want all datasets (startup memory/time increases) BATCH_SIZE = 16 TOP_K_DEFAULT = 3 # Hugging Face router endpoint (new inference endpoint) HF_API_URL = "https://router.huggingface.co/hf-inference" HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN", None) # Device - CPU only device = torch.device("cpu") print("Running on device:", device) # ------------------------- # Load dataset and SigLip model & precompute text embeddings at startup # ------------------------- print("Loading datasets and computing SigLip text embeddings (startup)...") texts_all: List[str] = [] for i in range(1, NUM_DATASETS + 1): ds = load_dataset(DATASET_TEMPLATE.format(i), split="train") texts_all.extend(ds["text"]) siglip_processor = AutoProcessor.from_pretrained(SIGLIP_MODEL_ID) siglip_model = AutoModel.from_pretrained(SIGLIP_MODEL_ID).to(device) siglip_model.eval() # Precompute text embeddings (on CPU) text_embeds_parts = [] for i in tqdm(range(0, len(texts_all), BATCH_SIZE), desc="Encoding texts (CPU)"): batch_texts = texts_all[i : i + BATCH_SIZE] inputs = siglip_processor(text=batch_texts, padding=True, truncation=True, return_tensors="pt") with torch.no_grad(): text_embeds = siglip_model.get_text_features(**inputs) text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) text_embeds_parts.append(text_embeds.cpu()) del inputs, text_embeds if text_embeds_parts: text_embeds_all = torch.cat(text_embeds_parts, dim=0) else: text_embeds_all = torch.empty((0, 0)) print(f"Encoded {len(texts_all)} texts. Embeddings shape: {text_embeds_all.shape}") # ------------------------- # Llava loading: try local package -> trust_remote_code -> HF Inference API (if token provided) # ------------------------- llava_tokenizer: Optional[AutoTokenizer] = None llava_model = None llava_mode: Optional[str] = None # 'local', 'trust_remote_code', 'hf_api', or None load_errors = [] # Attempt 1: local llava package (preferred) try: # this import requires the LLaVA repo to be installed in the environment (requirements.txt) from llava.model import LlavaForCausalLM # type: ignore print("Loading LlavaForCausalLM from installed 'llava' package (CPU)...") llava_tokenizer = AutoTokenizer.from_pretrained(LLAVA_MODEL_ID, use_fast=False) llava_model = LlavaForCausalLM.from_pretrained( LLAVA_MODEL_ID, device_map={"": "cpu"}, torch_dtype=torch.float32, low_cpu_mem_usage=True, ) llava_model.to(device) llava_model.eval() llava_mode = "local" print("✅ Llava loaded from installed package.") except Exception: tb_local = traceback.format_exc() load_errors.append(("local_llava_import", tb_local)) print("Local llava import failed — will try trust_remote_code fallback. See logs for details.") # Attempt 2: trust_remote_code fallback if llava_mode is None: try: print("Attempting AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True) (CPU)...") llava_tokenizer = AutoTokenizer.from_pretrained(LLAVA_MODEL_ID, use_fast=False) llava_model = AutoModelForCausalLM.from_pretrained( LLAVA_MODEL_ID, trust_remote_code=True, device_map={"": "cpu"}, torch_dtype=torch.float32, low_cpu_mem_usage=True, ) llava_model.to(device) llava_model.eval() llava_mode = "trust_remote_code" print("✅ Llava loaded via trust_remote_code fallback.") except Exception: tb_trust = traceback.format_exc() load_errors.append(("fallback_trust_remote_code", tb_trust)) print("trust_remote_code fallback failed — will try HF router if token provided.") # Attempt 3: Hugging Face router Inference API fallback (requires HUGGINGFACE_TOKEN) if llava_mode is None and HUGGINGFACE_TOKEN: llava_mode = "hf_api" print("No usable local model found. Will use Hugging Face router Inference API for generation (HUGGINGFACE_TOKEN detected).") if llava_mode is None: print("WARNING: No Llava model available and no HUGGINGFACE_TOKEN supplied. Generation will return an actionable error.") for name, tb in load_errors: print(f"--- {name} traceback ---\n{tb}") # ------------------------- # Helper: call Hugging Face router inference API # ------------------------- def call_hf_inference_api(prompt: str, max_new_tokens: int = 256, temperature: float = 0.0): if not HUGGINGFACE_TOKEN: raise RuntimeError("HUGGINGFACE_TOKEN not set; cannot call Hugging Face Inference API.") headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}", "Content-Type": "application/json"} payload = { "model": LLAVA_MODEL_ID, "inputs": prompt, "parameters": {"max_new_tokens": max_new_tokens, "temperature": temperature}, "options": {"wait_for_model": True}, } resp = requests.post(HF_API_URL, headers=headers, json=payload, timeout=300) if resp.status_code != 200: raise RuntimeError(f"HF Inference API error {resp.status_code}: {resp.text}") data = resp.json() # handle common response shapes if isinstance(data, list) and data and isinstance(data[0], dict) and "generated_text" in data[0]: return data[0]["generated_text"] if isinstance(data, dict) and "generated_text" in data: return data["generated_text"] if isinstance(data, str): return data return json.dumps(data) # ------------------------- # Retrieval & generation # ------------------------- def retrieve_top_k_texts(image: Image.Image, k: int = TOP_K_DEFAULT): inputs = siglip_processor(images=image, return_tensors="pt") with torch.no_grad(): img_embed = siglip_model.get_image_features(**inputs) img_embed = img_embed / img_embed.norm(p=2, dim=-1, keepdim=True) sims = F.cosine_similarity(img_embed.cpu(), text_embeds_all) topk = torch.topk(sims, k) results = [(texts_all[idx.item()], float(score)) for idx, score in zip(topk.indices, topk.values)] return results def llava_answer(image: Image.Image, retrieved_texts, question: str, max_tokens: int = 256): context_text = "\n".join([f"Retrieved Text: {t}" for t, _ in retrieved_texts]) prompt = ( "You are an agricultural assistant. Use the provided retrieved texts to answer concisely.\n\n" f"Retrieved texts:\n{context_text}\n\n" f"User question: {question}\n\n" "Provide a concise, actionable answer and crop suggestions when applicable." ) if llava_mode in ("local", "trust_remote_code"): inputs = llava_tokenizer(prompt, return_tensors="pt") inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): output_ids = llava_model.generate(**inputs, max_new_tokens=max_tokens) resp = llava_tokenizer.decode(output_ids[0], skip_special_tokens=True) return resp elif llava_mode == "hf_api": return call_hf_inference_api(prompt, max_new_tokens=max_tokens) else: err = ( "No Llava model is available for generation.\n\n" "Fix options:\n" "1) Install the LLaVA repo in requirements.txt and rebuild the Space:\n" " git+https://github.com/haotian-liu/LLaVA.git@main\n" "2) Or add a valid Hugging Face API token as HUGGINGFACE_TOKEN in Space secrets to use the router.\n\n" "Check Space logs for detailed tracebacks printed at startup." ) return err # ------------------------- # Gradio app # ------------------------- def gradio_pipeline(image: Image.Image, question: str, k: int = TOP_K_DEFAULT): if image is None or not question: return None, "Please provide both an image and a question." retrieved = retrieve_top_k_texts(image, k=int(k)) try: answer = llava_answer(image, retrieved, question) except Exception as e: tb = traceback.format_exc() answer = f"Error during generation: {e}\n\nTraceback:\n{tb}" return image, answer with gr.Blocks(title="Agri Image + Question → Llava Response (robust)") as demo: gr.Markdown( "## Agri Image QA\n\nThis app preloads SigLip embeddings at startup. " "Generation uses a local Llava model if available, otherwise the Hugging Face router Inference API " "(requires HUGGINGFACE_TOKEN secret in Space settings)." ) with gr.Row(): img_in = gr.Image(type="pil") out_img = gr.Image(type="pil", label="Image") question_input = gr.Textbox(label="Question about the image", lines=2) k_slider = gr.Slider(minimum=1, maximum=10, step=1, value=TOP_K_DEFAULT, label="Top-k retrieval") txt_out = gr.Textbox(label="Llava Response", lines=12) run_btn = gr.Button("Generate Answer") run_btn.click(fn=gradio_pipeline, inputs=[img_in, question_input, k_slider], outputs=[out_img, txt_out]) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", share=False)