app.py

app.py

@@ -1,13 +1,18 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForVision2Seq
-import torch
+from transformers import AutoProcessor, LlavaForConditionalGeneration
 from PIL import Image
+import torch
 
-# Load model and processor
 MODEL_ID = "liuhaotian/llava-v1.6-vicuna-7b"
-processor = AutoProcessor.from_pretrained(MODEL_ID)
-model = AutoModelForVision2Seq.from_pretrained(MODEL_ID, torch_dtype=torch.float16, low_cpu_mem_usage=True)
 
+# Load model and processor (use correct classes)
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = LlavaForConditionalGeneration.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True
+)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
@@ -23,37 +28,36 @@ def detect_object(image, prompt):
         return "⚠️ Please upload a base image first."
 
     query = (
-        f"Ignore the base image and focus only on new or added objects. "
-        f"Base image and detection image are given. {prompt or 'Describe and identify the materials of new objects.'}"
+        f"Ignore the base image and only analyze the differences. "
+        f"{prompt or 'Detect new objects and identify their material type.'}"
     )
 
     inputs = processor(
         text=query,
         images=[base_image, image],
         return_tensors="pt"
-    ).to(device, torch.float16)
+    ).to(device, torch.float16 if torch.cuda.is_available() else torch.float32)
 
     output = model.generate(**inputs, max_new_tokens=256)
     result = processor.decode(output[0], skip_special_tokens=True)
     return result
 
-# Build Gradio UI
-with gr.Blocks(title="LLaVA Object & Material Detector") as demo:
-    gr.Markdown("## 🧠 LLaVA-1.6 (Vicuna-7B) Object & Material Detection\nUpload a base image first, then upload another image to detect new objects while ignoring the base.")
+with gr.Blocks(title="LLaVA Object Detector") as demo:
+    gr.Markdown("## 🧠 LLaVA 1.6 Vicuna-7B — Visual Detection & Material Identification")
 
     with gr.Row():
         with gr.Column():
             base_img = gr.Image(label="Base Image", type="pil")
             set_base_btn = gr.Button("Set as Base Image")
-            set_base_status = gr.Textbox(label="Status")
+            base_status = gr.Textbox(label="Status")
 
         with gr.Column():
             target_img = gr.Image(label="Detection Image", type="pil")
-            user_prompt = gr.Textbox(label="Custom Instruction", placeholder="e.g. Detect added objects and describe their material.")
+            prompt = gr.Textbox(label="Instruction", placeholder="Detect new objects and describe material")
             run_btn = gr.Button("Run Detection")
-            output_box = gr.Textbox(label="Model Output")
+            output_box = gr.Textbox(label="Output")
 
-    set_base_btn.click(set_base, inputs=base_img, outputs=set_base_status)
-    run_btn.click(detect_object, inputs=[target_img, user_prompt], outputs=output_box)
+    set_base_btn.click(set_base, inputs=base_img, outputs=base_status)
+    run_btn.click(detect_object, inputs=[target_img, prompt], outputs=output_box)
 
 demo.launch()