Update app.py

app.py

@@ -1,60 +1,41 @@
-# app.py
 import torch
-from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
+from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration, BitsAndBytesConfig
 import gradio as gr
 from PIL import Image
 import re
 from typing import List, Tuple
-import os
-import logging
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
+# Configuration for 4-bit quantization
+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):
-        # Check if CUDA is available
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
-        logger.info(f"Using device: {self.device}, dtype: {torch_dtype}")
-
         try:
-            # Load processor
-            logger.info("Loading processor...")
+            # Initialize InstructBLIP-FLAN-T5-XL with 4-bit quantization
             self.processor = InstructBlipProcessor.from_pretrained(
-                "Salesforce/instructblip-vicuna-7b"
+                "Salesforce/instructblip-flan-t5-xl",
+                cache_dir="model_cache"
             )
-
-            # Load model with appropriate settings
-            logger.info("Loading model...")
             self.model = InstructBlipForConditionalGeneration.from_pretrained(
-                "Salesforce/instructblip-vicuna-7b",
+                "Salesforce/instructblip-flan-t5-xl",
+                quantization_config=quant_config,
                 device_map="auto",
-                torch_dtype=torch_dtype,
-                load_in_8bit=True if self.device == "cuda" else False,
-                offload_folder="./offload",
+                torch_dtype=torch.float16,
+                cache_dir="model_cache"
             )
-            logger.info("Model loaded successfully")
-
         except Exception as e:
-            logger.error(f"Error loading model: {str(e)}")
-            raise RuntimeError(f"Failed to initialize model: {str(e)}")
+            raise RuntimeError(f"Model loading failed: {str(e)}")
 
         self.pollutants = [
-            "plastic waste",
-            "chemical foam",
-            "industrial discharge",
-            "sewage water",
-            "oil spill",
-            "organic debris",
-            "construction waste",
-            "medical waste",
-            "floating trash",
-            "algal bloom",
-            "toxic sludge",
-            "agricultural runoff",
+            "plastic waste", "chemical foam", "industrial discharge",
+            "sewage water", "oil spill", "organic debris",
+            "construction waste", "medical waste", "floating trash",
+            "algal bloom", "toxic sludge", "agricultural runoff"
         ]
 
         self.severity_descriptions = {
@@ -67,16 +48,15 @@ class RiverPollutionAnalyzer:
             7: "Very severe pollution - Major ecosystem impact",
             8: "Extreme pollution - Dangerous levels",
             9: "Critical pollution - Immediate action needed",
-            10: "Disaster level - Ecological catastrophe",
+            10: "Disaster level - Ecological catastrophe"
         }
 
     def analyze_image(self, image):
         """Analyze river pollution with robust parsing"""
-        try:
-            if not isinstance(image, Image.Image):
-                image = Image.fromarray(image)
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
 
-            prompt = """Analyze this river pollution scene and provide:
+        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
 
@@ -84,9 +64,12 @@ Respond EXACTLY in this format:
 Pollutants: [comma separated list]
 Severity: [number]"""
 
-            inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(
-                self.device
-            )
+        try:
+            inputs = self.processor(
+                images=image,
+                text=prompt,
+                return_tensors="pt"
+            ).to(self.model.device)
 
             with torch.no_grad():
                 outputs = self.model.generate(
@@ -94,38 +77,39 @@ Severity: [number]"""
                     max_new_tokens=200,
                     temperature=0.5,
                     top_p=0.85,
-                    do_sample=True,
+                    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:
-            logger.error(f"Error analyzing image: {str(e)}")
-            return "⚠️ Error analyzing image. Please try again or check the logs for details."
+            return f"⚠️ Analysis failed: {str(e)}"
 
+    # [Keep your existing parsing/formatting methods]
     def _parse_response(self, analysis: str) -> Tuple[List[str], int]:
-        """Robust parsing of model response"""
         pollutants = []
         severity = 3
 
         # Extract pollutants
         pollutant_match = re.search(
-            r"(?i)(pollutants?|contaminants?)[:\s]*\[?(.*?)(?:\]|Severity|severity|$)",
-            analysis,
+            r'(?i)(pollutants?|contaminants?)[:\s]*\[?(.*?)(?:\]|Severity|severity|$)',
+            analysis
         )
 
         if pollutant_match:
             pollutants_str = pollutant_match.group(2).strip()
             pollutants = [
                 p.strip().lower()
-                for p in re.split(r"[,;]|\band\b", pollutants_str)
+                for p in re.split(r'[,;]|\band\b', pollutants_str)
                 if p.strip().lower() in self.pollutants
             ]
 
         # Extract severity
-        severity_match = re.search(r"(?i)(severity|level)[:\s]*(\d{1,2})", analysis)
+        severity_match = re.search(
+            r'(?i)(severity|level)[:\s]*(\d{1,2})',
+            analysis
+        )
 
         if severity_match:
             try:
@@ -138,72 +122,46 @@ Severity: [number]"""
         return pollutants, severity
 
     def _calculate_severity(self, pollutants: List[str]) -> int:
-        """Weighted severity calculation"""
         if not pollutants:
             return 1
 
         weights = {
-            "medical waste": 3,
-            "toxic sludge": 3,
-            "oil spill": 2.5,
-            "chemical foam": 2,
-            "industrial discharge": 2,
-            "sewage water": 2,
-            "plastic waste": 1.5,
-            "construction waste": 1.5,
-            "algal bloom": 1.5,
-            "agricultural runoff": 1.5,
-            "floating trash": 1,
-            "organic debris": 1,
+            "medical waste": 3, "toxic sludge": 3, "oil spill": 2.5,
+            "chemical foam": 2, "industrial discharge": 2, "sewage water": 2,
+            "plastic waste": 1.5, "construction waste": 1.5, "algal bloom": 1.5,
+            "agricultural runoff": 1.5, "floating trash": 1, "organic debris": 1
         }
 
         avg_weight = sum(weights.get(p, 1) for p in pollutants) / len(pollutants)
         return min(10, max(1, round(avg_weight * 3)))
 
     def _format_analysis(self, pollutants: List[str], severity: int) -> str:
-        """Generate formatted report"""
         severity_bar = f"""📊 Severity: {severity}/10
 {"█" * severity}{"░" * (10 - severity)}
-{self.severity_descriptions.get(severity, "")}"""
+{self.severity_descriptions.get(severity, '')}"""
 
-        pollutants_list = (
-            "\n🔍 No pollutants detected"
-            if not pollutants
-            else "\n".join(
-                f"{i}. {p.capitalize()}" for i, p in enumerate(pollutants[:5], 1)
-            )
-        )
+        pollutants_list = "\n🔍 No pollutants detected" if not pollutants else "\n".join(
+            f"{i}. {p.capitalize()}" for i, p in enumerate(pollutants[:5], 1))
 
         return f"""🌊 River Pollution Analysis 🌊
 {pollutants_list}
 {severity_bar}"""
 
     def analyze_chat(self, message: str) -> str:
-        """Handle chat questions about pollution"""
-        try:
-            message = message.lower().strip()
-            if any(word in message for word in ["hello", "hi", "hey"]):
-                return "Hello! I'm a river pollution analyzer. Ask me about pollution types or upload an image for analysis."
-            elif "pollution" in message:
-                return "Common river pollutants include: plastic waste, chemical foam, industrial discharge, sewage water, and oil spills."
-            elif "severity" in message:
-                return "Severity is rated 1-10 (1=minimal, 10=disaster). It considers pollutant type and quantity."
-            elif "help" in message:
-                return "I can: 1) Analyze river pollution in images 2) Answer pollution questions. Try uploading an image or asking about pollution types."
-            else:
-                return "I specialize in river pollution analysis. Try asking about pollution types or upload an image for analysis."
-        except Exception as e:
-            logger.error(f"Error in chat: {str(e)}")
-            return "Sorry, I encountered an error. Please try again."
-
+        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 or upload an image for analysis."
+        elif "pollution" in message.lower():
+            return "Common river pollutants include: plastic waste, chemical foam, industrial discharge, sewage water, and oil spills."
+        else:
+            return "I can answer questions about river pollution. Try asking about pollution types or upload an image for analysis."
 
-# Initialize analyzer
+# Initialize with error handling
 try:
     analyzer = RiverPollutionAnalyzer()
-    logger.info("Analyzer initialized successfully")
+    model_status = "✅ Model loaded successfully"
 except Exception as e:
-    logger.error(f"Failed to initialize analyzer: {str(e)}")
-    raise
+    analyzer = None
+    model_status = f"❌ Model loading failed: {str(e)}"
 
 css = """
 .header {
@@ -213,16 +171,13 @@ css = """
     border-radius: 10px;
     margin-bottom: 20px;
 }
-
 .side-by-side {
     display: flex;
     gap: 20px;
 }
-
 .left-panel, .right-panel {
     flex: 1;
 }
-
 .analysis-box {
     padding: 20px;
     background: #f8f9fa;
@@ -230,97 +185,72 @@ css = """
     margin-top: 20px;
     border: 1px solid #dee2e6;
 }
-
 .chat-container {
     background: #f8f9fa;
     padding: 20px;
     border-radius: 10px;
     height: 100%;
 }
-
-.pollution-icon {
-    font-size: 24px;
-    margin-right: 10px;
-}
-
-.severity-bar {
-    font-family: monospace;
-    font-size: 16px;
-}
-
-.error-message {
-    color: #dc3545;
-    font-weight: bold;
-}
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     with gr.Column(elem_classes="header"):
         gr.Markdown("# 🌍 River Pollution Analyzer")
-        gr.Markdown("### AI-powered water pollution detection")
+        gr.Markdown(f"### {model_status}")
 
     with gr.Row(elem_classes="side-by-side"):
-        # Left Panel
         with gr.Column(elem_classes="left-panel"):
             with gr.Group():
-                image_input = gr.Image(
-                    type="pil", label="Upload River Image", height=300
-                )
+                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")
                 analysis_output = gr.Markdown()
 
-        # Right Panel
         with gr.Column(elem_classes="right-panel"):
             with gr.Group(elem_classes="chat-container"):
-                chatbot = gr.Chatbot(
-                    label="Pollution Analysis Q&A", height=400, bubble_full_width=False
-                )
+                chatbot = gr.Chatbot(label="Pollution Analysis Q&A", height=400)
                 with gr.Row():
                     chat_input = gr.Textbox(
                         placeholder="Ask about pollution sources...",
                         label="Your Question",
                         container=False,
-                        scale=5,
+                        scale=5
                     )
                     chat_btn = gr.Button("💬 Ask", variant="secondary", scale=1)
                 clear_btn = gr.Button("🧹 Clear Chat History", size="sm")
 
     analyze_btn.click(
-        analyzer.analyze_image, inputs=image_input, outputs=analysis_output
+        analyzer.analyze_image if analyzer else lambda x: "Model not loaded",
+        inputs=image_input,
+        outputs=analysis_output
     )
 
     chat_input.submit(
         lambda msg, chat: ("", chat + [(msg, analyzer.analyze_chat(msg))]),
         inputs=[chat_input, chatbot],
-        outputs=[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],
+        outputs=[chat_input, chatbot]
     )
 
     clear_btn.click(lambda: None, outputs=[chatbot])
 
     gr.Examples(
         examples=[
-            [
-                "https://huggingface.co/spaces/atharwaah1work/tarak.AI/resolve/main/polluted_river1.jpg"
-            ],
-            [
-                "https://huggingface.co/spaces/atharwaah1work/tarak.AI/resolve/main/polluted_river2.jpg"
-            ],
+            ["https://huggingface.co/spaces/atharwaah1work/tarak.AI/resolve/main/polluted_river1.jpg"],
+            ["https://huggingface.co/spaces/atharwaah1work/tarak.AI/resolve/main/polluted_river2.jpg"]
         ],
         inputs=image_input,
         outputs=analysis_output,
-        fn=analyzer.analyze_image,
+        fn=analyzer.analyze_image if analyzer else lambda x: "Model not loaded",
         cache_examples=True,
-        label="Try example images:",
+        label="Try example images:"
     )
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.queue(max_size=3).launch()