Update app.py

app.py

@@ -1,144 +1,86 @@
-# app.py
+# app.py - River Pollution Analyzer with instructblip-flan-t5-xl
 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
-
-# Set cache directory (important for Hugging Face Spaces)
-CACHE_DIR = "model_cache"
-os.makedirs(CACHE_DIR, exist_ok=True)
-
-
-class RiverPollutionAnalyzer:
-    def __init__(self):
-        try:
-            # Load with caching and optimized device placement
-            self.processor = InstructBlipProcessor.from_pretrained(
-                "Salesforce/instructblip-vicuna-7b", cache_dir=CACHE_DIR
-            )
-
-            self.model = InstructBlipForConditionalGeneration.from_pretrained(
-                "Salesforce/instructblip-vicuna-7b",
-                cache_dir=CACHE_DIR,
-                device_map="auto",
-                torch_dtype=torch.float16,
-                offload_folder="offload",
-                low_cpu_mem_usage=True,
-            )
-
-        except Exception as 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",
-        ]
-
-        self.severity_descriptions = {
-            1: "Minimal pollution - Slightly noticeable",
-            2: "Minor pollution - Small amounts visible",
-            3: "Moderate pollution - Clearly visible",
-            4: "Significant pollution - Affecting water quality",
-            5: "Heavy pollution - Obvious environmental impact",
-            6: "Severe pollution - Large accumulation",
-            7: "Very severe pollution - Major ecosystem impact",
-            8: "Extreme pollution - Dangerous levels",
-            9: "Critical pollution - Immediate action needed",
-            10: "Disaster level - Ecological catastrophe",
-        }
-
-    def analyze_image(self, image):
-        """Analyze river pollution with robust parsing"""
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-
-        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]"""
-
-        inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(
-            self.model.device
+import logging
+import functools
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@functools.cache
+def get_analyzer():
+    logger.info("Loading instructblip-flan-t5-xl...")
+    try:
+        # 4-bit config (works on GPU if 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,
         )
 
+        processor = InstructBlipProcessor.from_pretrained("Salesforce/instructblip-flan-t5-xl")
+        
+        model = InstructBlipForConditionalGeneration.from_pretrained(
+            "Salesforce/instructblip-flan-t5-xl",
+            quantization_config=quant_config if torch.cuda.is_available() else None,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto"
+        )
+        return processor, model
+        
+    except Exception as e:
+        logger.error(f"Model load failed: {str(e)}")
+        raise RuntimeError("Model loading error. Check logs.")
+
+def analyze_image(image):
+    try:
+        processor, model = get_analyzer()
+        prompt = """Analyze river pollution. List pollutants and severity (1-10).
+Respond EXACTLY like this:
+Pollutants: [list]
+Severity: [number]"""
+        
+        inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
+        
         with torch.no_grad():
-            outputs = self.model.generate(
+            outputs = model.generate(
                 **inputs,
-                max_new_tokens=200,
-                temperature=0.5,
-                top_p=0.85,
-                do_sample=True,
+                max_new_tokens=100,
+                temperature=0.7
             )
-
-        analysis = self.processor.batch_decode(outputs, skip_special_tokens=True)[0]
-        pollutants, severity = self._parse_response(analysis)
-        return self._format_analysis(pollutants, severity)
-
-    # ... [keep all your other methods unchanged] ...
-
-
-# Initialize analyzer with error handling
-try:
-    analyzer = RiverPollutionAnalyzer()
-    model_status = "Model loaded successfully!"
-except Exception as e:
-    analyzer = None
-    model_status = f"Model failed to load: {str(e)}"
-    print(model_status)
-
-
-# Create wrapper function for Gradio
-def analyze_image_wrapper(image):
-    if analyzer is None:
-        return (
-            f"⚠️ Error: {model_status}\nPlease try again later or use a smaller image."
-        )
-    return analyzer.analyze_image(image)
-
-
-# Gradio interface
-css = """
-/* [keep your existing CSS] */
-"""
-
-with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    # [keep all your existing UI code]
-
-    analyze_btn.click(
-        analyze_image_wrapper, inputs=image_input, outputs=analysis_output
-    )
-
-    # Update examples to use wrapper
-    gr.Examples(
-        examples=[
-            [
-                "https://huggingface.co/spaces/Atharwaaah/SLCR-FLOWCODE-tarak.AI/resolve/main/polluted_river1.jpg"
-            ],
-            [
-                "https://huggingface.co/spaces/Atharwaaah/SLCR-FLOWCODE-tarak.AI/resolve/main/polluted_river2.jpg"
-            ],
-        ],
-        inputs=image_input,
-        outputs=analysis_output,
-        fn=analyze_image_wrapper,
-        cache_examples=True,
-        label="Try example images:",
-    )
-
-demo.launch()
+            
+        result = processor.decode(outputs[0], skip_special_tokens=True)
+        
+        # Format output
+        if "Pollutants:" in result and "Severity:" in result:
+            pollutants = result.split("Pollutants:")[1].split("Severity:")[0].strip()
+            severity = result.split("Severity:")[1].strip()
+            return f"""🌊 Analysis Result:
+📌 Pollutants: {pollutants}
+📈 Severity: {severity}/10"""
+        return result
+        
+    except Exception as e:
+        logger.error(f"Error: {str(e)}")
+        return f"⚠️ Error (try a smaller image): {str(e)}"
+
+# Minimal UI
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🌍 River Pollution Analyzer (instructblip-flan-t5-xl)")
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload Image")
+        analyze_btn = gr.Button("Analyze", variant="primary")
+    output = gr.Textbox(label="Result")
+    
+    analyze_btn.click(analyze_image, inputs=image_input, outputs=output)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)