Update app.py

app.py

@@ -1,34 +1,53 @@
-# app.py - CPU-Compatible Version for Hugging Face Spaces
+!pip install -q transformers accelerate bitsandbytes gradio torch pillow
+
 import torch
-from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
+from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration, BitsAndBytesConfig
 import gradio as gr
 from PIL import Image
 import re
 import os
 from typing import List, Tuple
 
-# Create model cache directory
-os.makedirs("model_cache", exist_ok=True)
+# Configuration for 4-bit quantization (if GPU available)
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True
+)
 
 class RiverPollutionAnalyzer:
     def __init__(self):
         try:
-            # Initialize model (CPU-compatible version)
+            # Initialize model with fallback for CPU
             self.processor = InstructBlipProcessor.from_pretrained(
                 "Salesforce/instructblip-flan-t5-xl",
                 cache_dir="model_cache"
             )
-            self.model = InstructBlipForConditionalGeneration.from_pretrained(
-                "Salesforce/instructblip-flan-t5-xl",
-                device_map="auto",
-                torch_dtype=torch.float32,  # Using float32 instead of quantization
-                cache_dir="model_cache",
-                low_cpu_mem_usage=True
-            )
-            self.model_loaded = True
-            self.status = "✅ Model loaded successfully (CPU mode)"
+            
+            if torch.cuda.is_available():
+                self.model = InstructBlipForConditionalGeneration.from_pretrained(
+                    "Salesforce/instructblip-flan-t5-xl",
+                    device_map="auto",
+                    quantization_config=quant_config,
+                    torch_dtype=torch.float16,
+                    cache_dir="model_cache"
+                )
+                self.device = "cuda"
+                self.status = "✅ Model loaded (4-bit GPU)"
+            else:
+                self.model = InstructBlipForConditionalGeneration.from_pretrained(
+                    "Salesforce/instructblip-flan-t5-xl",
+                    device_map="auto",
+                    torch_dtype=torch.float32,
+                    cache_dir="model_cache",
+                    low_cpu_mem_usage=True
+                )
+                self.device = "cpu"
+                self.status = "⚠️ Model loaded (CPU mode - slower)"
+                
         except Exception as e:
-            self.model_loaded = False
+            self.model = None
             self.status = f"❌ Model loading failed: {str(e)}"
             print(self.status)
 
@@ -53,47 +72,49 @@ class RiverPollutionAnalyzer:
         }
 
     def analyze_image(self, image):
-        """Analyze river pollution with CPU optimizations"""
-        if not self.model_loaded:
+        """Analyze river pollution with device-aware processing"""
+        if not self.model:
             return "Model not loaded. Please check logs."
             
-        try:
-            if not isinstance(image, Image.Image):
-                image = Image.fromarray(image)
-            
-            # Resize for CPU efficiency
-            image = image.resize((512, 512))
-            
-            prompt = """Analyze this river pollution and list:
-1. Visible pollutants from: [plastic waste, chemical foam, industrial discharge, sewage water, oil spill, organic debris, construction waste, medical waste, floating trash, algal bloom, toxic sludge, agricultural runoff]
-2. Severity estimate (1-10)
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+        
+        # Resize for efficiency
+        image = image.resize((512, 512))
 
-Respond EXACTLY as:
+        prompt = """Analyze this river pollution scene and provide:
+1. List ALL visible pollutants ONLY from: [plastic waste, chemical foam, industrial discharge, sewage water, oil spill, organic debris, construction waste, medical waste, floating trash, algal bloom, toxic sludge, agricultural runoff]
+2. Estimate pollution severity from 1-10
+
+Respond EXACTLY in this format:
 Pollutants: [comma separated list]
 Severity: [number]"""
 
+        try:
             inputs = self.processor(
                 images=image,
                 text=prompt,
                 return_tensors="pt"
-            )  # No .to("cuda") for CPU
+            ).to(self.model.device)
 
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=100,  # Reduced for CPU
-                temperature=0.5,
-                top_p=0.85
-            )
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=150,  # Reduced for stability
+                    temperature=0.5,
+                    top_p=0.85,
+                    do_sample=True
+                )
 
             analysis = self.processor.batch_decode(outputs, skip_special_tokens=True)[0]
             pollutants, severity = self._parse_response(analysis)
             return self._format_analysis(pollutants, severity)
-            
         except Exception as e:
             return f"⚠️ Analysis error: {str(e)}"
 
+    # [Keep all existing helper methods unchanged]
     def _parse_response(self, analysis: str) -> Tuple[List[str], int]:
-        """Parse model response"""
+        """Same parsing logic as before"""
         pollutants = []
         severity = 3
 
@@ -123,7 +144,7 @@ Severity: [number]"""
         return pollutants, severity
 
     def _calculate_severity(self, pollutants: List[str]) -> int:
-        """Calculate weighted severity"""
+        """Same severity calculation"""
         if not pollutants:
             return 1
 
@@ -138,7 +159,7 @@ Severity: [number]"""
         return min(10, max(1, round(avg_weight * 3)))
 
     def _format_analysis(self, pollutants: List[str], severity: int) -> str:
-        """Generate formatted report"""
+        """Same formatting"""
         severity_bar = f"""📊 Severity: {severity}/10
 {"█" * severity}{"░" * (10 - severity)}
 {self.severity_descriptions.get(severity, '')}"""
@@ -150,6 +171,15 @@ Severity: [number]"""
 {pollutants_list}
 {severity_bar}"""
 
+    def analyze_chat(self, message: str) -> str:
+        """Handle chat questions"""
+        if any(word in message.lower() for word in ["hello", "hi", "hey"]):
+            return "Hello! I'm a river pollution analyzer. Ask me about pollution types."
+        elif "pollution" in message.lower():
+            return "Common river pollutants: plastic waste, chemical foam, industrial discharge, sewage water, oil spills."
+        else:
+            return "I can answer questions about river pollution. Try asking about pollution types."
+
 # Initialize analyzer
 analyzer = RiverPollutionAnalyzer()
 
@@ -186,9 +216,6 @@ css = """
     background: #2a2a2a;
     border-color: #444;
 }
-.btn-primary {
-    margin-top: 10px;
-}
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
@@ -200,15 +227,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         # Left Panel
         with gr.Column(elem_classes="left-panel"):
             with gr.Group():
-                image_input = gr.Image(
-                    type="pil",
-                    label="Upload River Image",
-                    height=300
-                )
-                analyze_btn = gr.Button(
-                    "🔍 Analyze Pollution",
-                    variant="primary"
-                )
+                image_input = gr.Image(type="pil", label="Upload River Image", height=300)
+                analyze_btn = gr.Button("🔍 Analyze Pollution", variant="primary")
             
             with gr.Group(elem_classes="analysis-box"):
                 gr.Markdown("### 📊 Analysis Report")
@@ -217,16 +237,16 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         # Right Panel
         with gr.Column(elem_classes="right-panel"):
             with gr.Group(elem_classes="chat-container"):
-                gr.Markdown("### 💬 Pollution Q&A")
-                chatbot = gr.Chatbot(height=400)
+                chatbot = gr.Chatbot(label="Pollution Q&A", height=400)
                 with gr.Row():
                     chat_input = gr.Textbox(
-                        placeholder="Ask about pollution types...",
+                        placeholder="Ask about pollution sources...",
+                        label="Your Question",
                         container=False,
                         scale=5
                     )
-                    chat_btn = gr.Button("Send", variant="secondary", scale=1)
-                clear_btn = gr.Button("Clear Chat", size="sm")
+                    chat_btn = gr.Button("💬 Ask", variant="secondary", scale=1)
+                clear_btn = gr.Button("🧹 Clear Chat", size="sm")
 
     analyze_btn.click(
         analyzer.analyze_image,
@@ -234,7 +254,20 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=analysis_output
     )
 
-    # Example images (host them in your Space's repository)
+    chat_input.submit(
+        lambda msg, chat: ("", chat + [(msg, analyzer.analyze_chat(msg))]),
+        inputs=[chat_input, chatbot],
+        outputs=[chat_input, chatbot]
+    )
+
+    chat_btn.click(
+        lambda msg, chat: ("", chat + [(msg, analyzer.analyze_chat(msg))]),
+        inputs=[chat_input, chatbot],
+        outputs=[chat_input, chatbot]
+    )
+
+    clear_btn.click(lambda: None, outputs=[chatbot])
+
     gr.Examples(
         examples=[
             ["examples/polluted_river1.jpg"],
@@ -243,7 +276,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         inputs=image_input,
         outputs=analysis_output,
         fn=analyzer.analyze_image,
-        cache_examples=False,  # Disabled for CPU
+        cache_examples=torch.cuda.is_available(),  # Cache only if GPU available
         label="Example Images"
     )