app.py

app.py

@@ -1,63 +1,38 @@
-import gradio as gr
-from transformers import AutoProcessor, LlavaForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForVision2Seq
 from PIL import Image
 import torch
+import gradio as gr
 
-MODEL_ID = "liuhaotian/llava-v1.6-vicuna-7b"
-
-# 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)
-
-base_image = None
-
-def set_base(image):
-    global base_image
-    base_image = image
-    return "✅ Base image has been set."
-
-def detect_object(image, prompt):
-    if base_image is None:
-        return "⚠️ Please upload a base image first."
-
-    query = (
-        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 if torch.cuda.is_available() else torch.float32)
-
-    output = model.generate(**inputs, max_new_tokens=256)
+MODEL_ID = "HuggingFaceM4/idefics2-8b"
+
+# Load model and processor
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = AutoModelForVision2Seq.from_pretrained(MODEL_ID, torch_dtype=torch.float16, device_map="auto")
+
+def analyze_images(base_img, target_img, user_prompt):
+    if base_img is None or target_img is None:
+        return "Please upload both a base image and a target image."
+    
+    images = [base_img, target_img]
+    prompt = f"Ignore the first image (base image). Analyze the second image: {user_prompt}"
+    
+    inputs = processor(images=images, text=prompt, return_tensors="pt").to(model.device)
+    output = model.generate(**inputs, max_new_tokens=200)
+    
     result = processor.decode(output[0], skip_special_tokens=True)
     return result
 
-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")
-            base_status = gr.Textbox(label="Status")
-
-        with gr.Column():
-            target_img = gr.Image(label="Detection Image", type="pil")
-            prompt = gr.Textbox(label="Instruction", placeholder="Detect new objects and describe material")
-            run_btn = gr.Button("Run Detection")
-            output_box = gr.Textbox(label="Output")
-
-    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 = gr.Interface(
+    fn=analyze_images,
+    inputs=[
+        gr.Image(type="pil", label="Base Image"),
+        gr.Image(type="pil", label="Target Image"),
+        gr.Textbox(label="Prompt", placeholder="Describe what to analyze...")
+    ],
+    outputs=gr.Textbox(label="Model Output"),
+    title="Image Comparison with IDEFICS2-8B",
+    description="Upload two images. The model will ignore the base image and analyze the target image according to your prompt."
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()