Update app.py

app.py

@@ -1,246 +1,205 @@
-import spaces
+# app.py
+# Dermatology-AI-Assistant — Hugging Face Space (ZeroGPU-ready)
+# - Logging is configured before use
+# - No runtime pip installs (use requirements.txt)
+# - ZeroGPU acquired only during inference via @spaces.GPU
+# - Uses qwen-vl-utils.process_vision_info (fixes missing attribute error)
+# - SSR disabled in Gradio launch to avoid Node 20 requirement in container
+
+import os
+import sys
+import logging
+from typing import Optional, Tuple
+
 import gradio as gr
+import spaces
 import torch
-from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
 from PIL import Image
-import logging
-import subprocess
-import sys
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
+from qwen_vl_utils import process_vision_info
 
-logging.basicConfig(level=logging.INFO)
+
+# ---------------------------
+# Logging
+# ---------------------------
+logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
 logger = logging.getLogger(__name__)
 
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# ---------------------------
+# Config
+# ---------------------------
+# Fine-tuned (or partially fine-tuned) Qwen VL checkpoint
+MODEL_ID = os.environ.get("MODEL_ID", "ColdSlim/Dermatology-Qwen2.5-VL-3B")
 
-# Global variables for model and processor
-model = None
-processor = None
+# Generation params (tweak as needed)
+GEN_KW = dict(
+    max_new_tokens=512,
+    do_sample=True,
+    temperature=0.7,
+    top_p=0.9,
+)
+
+# ZeroGPU time (seconds). Increase if your model is slow to generate.
+ZGPU_DURATION = int(os.environ.get("ZGPU_DURATION", "180"))
+
+# Preload only the processor on CPU; load the model inside GPU-decorated call.
+logger.info(f"Loading processor from: {MODEL_ID}")
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+logger.info("Processor loaded.")
+
+
+# ---------------------------
+# Helpers
+# ---------------------------
+def build_inputs(image: Image.Image, question: str):
+    """
+    Build Qwen-style multimodal chat inputs using qwen-vl-utils.
+    Returns a dict of tensors ready for model.generate.
+    """
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image},
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+
+    # Chat template
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    # Vision inputs
+    image_inputs, video_inputs = process_vision_info(messages)
+
+    # Pack tensors (CPU for now; we move to CUDA later)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+    return inputs
+
+
+def format_derm_disclaimer(ans: str) -> str:
+    """Append a short medical disclaimer (non-blocking)."""
+    tail = (
+        "\n\n---\n"
+        "_Disclaimer: This AI is not a medical device. The output is informational and may be inaccurate. "
+        "Consult a qualified dermatologist for diagnosis and treatment._"
+    )
+    return ans + tail
+
+
+# ---------------------------
+# Inference (ZeroGPU)
+# ---------------------------
+@spaces.GPU(duration=ZGPU_DURATION)
+def analyze_skin_condition(image: Optional[Image.Image], question: str) -> str:
+    """
+    Main inference function. Runs inside a ZeroGPU reservation window.
+    Loads model on GPU, generates, frees VRAM.
+    """
+    if image is None:
+        return "❌ Please upload an image first."
 
-def load_model():
-    """Load the fine-tuned dermatology model"""
-    global model, processor
-    
     try:
-        # Load the merged model (replace with your actual model path)
-        model_name = "ColdSlim/Dermatology-Qwen2.5-VL-3B"  # Update with your actual model name
-        
-        logger.info(f"Loading model: {model_name}")
-        processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+        logger.info(f"Loading model on GPU: {MODEL_ID}")
+        # On ZeroGPU, load inside the GPU-decorated function
         model = Qwen2VLForConditionalGeneration.from_pretrained(
-            model_name, 
-            dtype=torch.bfloat16,
-            device_map="auto",
+            MODEL_ID,
+            torch_dtype=torch.float16,          # fp16 is broadly compatible on ZeroGPU
+            device_map="cuda",                  # place modules on available CUDA
             trust_remote_code=True,
             low_cpu_mem_usage=True,
-            ignore_mismatched_sizes=True
+            ignore_mismatched_sizes=True,       # your logs indicated shape diffs; keep this to avoid crash
         )
-        
         logger.info("Model loaded successfully!")
-        return True
-        
-    except Exception as e:
-        logger.error(f"Error loading model: {e}")
-        return False
-
-def analyze_skin_condition(image, question="Describe this skin condition in detail."):
-    """Analyze skin condition from uploaded image"""
-    global model, processor
-    
-    if model is None or processor is None:
-        return "❌ Model not loaded. Please wait for the model to load or contact the administrator."
-    
-    if image is None:
-        return "❌ Please upload an image first."
-    
-    try:
-        # Prepare the conversation
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": image},
-                    {"type": "text", "text": question}
-                ]
-            }
-        ]
-        
-        # Process the input
-        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        image_inputs, video_inputs = processor.process_vision_info(messages)
-        
-        inputs = processor(
-            text=[text],
-            images=image_inputs,
-            videos=video_inputs,
-            padding=True,
-            return_tensors="pt"
-        )
-        
-        # Move inputs to the same device as model
-        inputs = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
-        
-        # Generate response
+
+        # Build and move inputs to CUDA
+        inputs = build_inputs(image, question)
+        inputs = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+
+        # Generate
         with torch.no_grad():
-            generated_ids = model.generate(
+            out_ids = model.generate(
                 **inputs,
-                max_new_tokens=512,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
-                pad_token_id=processor.tokenizer.eos_token_id
+                **GEN_KW,
+                pad_token_id=processor.tokenizer.eos_token_id,
             )
-        
-        # Decode the response
-        generated_ids_trimmed = [
-            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+
+        # Strip prompt tokens before decoding for clean answer
+        prompt_len_trimmed = [
+            out[len(inp):] for inp, out in zip(inputs["input_ids"], out_ids)
         ]
-        output_text = processor.batch_decode(
-            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        text = processor.batch_decode(
+            prompt_len_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
         )[0]
-        
-        return output_text
-        
+
+        # Free VRAM early
+        del model
+        torch.cuda.empty_cache()
+
+        return format_derm_disclaimer(text)
+
     except Exception as e:
-        logger.error(f"Error during inference: {e}")
-        return f"❌ Error analyzing image: {str(e)}"
-
-def create_interface():
-    """Create the Gradio interface"""
-    
-    # Load model on startup
-    model_loaded = load_model()
-    
-    with gr.Blocks(
-        title="Dermatology AI Assistant",
-        theme=gr.themes.Soft(),
-        css="""
-        .gradio-container {
-            max-width: 1200px !important;
-            margin: auto !important;
-        }
-        .main-header {
-            text-align: center;
-            margin-bottom: 2rem;
-        }
-        .warning-box {
-            background-color: #fff3cd;
-            border: 1px solid #ffeaa7;
-            border-radius: 8px;
-            padding: 1rem;
-            margin: 1rem 0;
-        }
-        """
-    ) as demo:
-        
-        gr.HTML("""
-        <div class="main-header">
-            <h1>🩺 Dermatology AI Assistant</h1>
-            <p>Powered by Qwen2.5-VL-3B fine-tuned for dermatology analysis</p>
-        </div>
-        """)
-        
-        # Warning message
-        gr.HTML("""
-        <div class="warning-box">
-            <h3>⚠️ Medical Disclaimer</h3>
-            <p>This AI assistant is for educational and research purposes only. 
-            It should not be used as a substitute for professional medical advice, 
-            diagnosis, or treatment. Always consult with a qualified healthcare 
-            provider for medical concerns.</p>
-        </div>
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                # Image upload
-                image_input = gr.Image(
-                    label="Upload Skin Image",
-                    type="pil",
-                    height=400
-                )
-                
-                # Question input
-                question_input = gr.Textbox(
-                    label="Question (Optional)",
-                    placeholder="Describe this skin condition in detail.",
-                    value="Describe this skin condition in detail.",
-                    lines=3
-                )
-                
-                # Analyze button
-                analyze_btn = gr.Button(
-                    "🔍 Analyze Skin Condition",
-                    variant="primary",
-                    size="lg"
-                )
-                
-                # Example questions
-                gr.HTML("""
-                <h4>💡 Example Questions:</h4>
-                <ul>
-                    <li>What type of skin condition is this?</li>
-                    <li>Describe the characteristics of this lesion.</li>
-                    <li>What are the potential causes of this skin issue?</li>
-                    <li>What should I know about this skin condition?</li>
-                </ul>
-                """)
-            
-            with gr.Column(scale=1):
-                # Output
-                output_text = gr.Textbox(
-                    label="AI Analysis",
-                    lines=15,
-                    max_lines=20,
-                    show_copy_button=True
-                )
-        
-        # Examples
-        gr.Examples(
-            examples=[
-                ["What type of skin condition is this?", "Describe this skin condition in detail."],
-                ["What are the characteristics of this lesion?", "Describe this skin condition in detail."],
-                ["What should I know about this skin issue?", "Describe this skin condition in detail."],
-            ],
-            inputs=[question_input, question_input],
-            label="💡 Example Questions"
+        logger.exception("Error during inference")
+        return f"❌ Error analyzing image: {e}"
+
+
+# ---------------------------
+# UI
+# ---------------------------
+def create_interface() -> gr.Blocks:
+    with gr.Blocks(title="Dermatology AI Assistant") as demo:
+        gr.Markdown(
+            "# Dermatology AI Assistant\n"
+            "Upload a skin photo and ask a question. The model will provide an informational response."
         )
-        
-        # Event handlers
-        analyze_btn.click(
-            fn=analyze_skin_condition,
-            inputs=[image_input, question_input],
-            outputs=output_text
+
+        with gr.Row():
+            image_input = gr.Image(type="pil", label="Upload Image (JPG/PNG)")
+            question_input = gr.Textbox(
+                label="Question / Prompt",
+                value="Describe this skin condition in detail and suggest possible next steps.",
+                lines=3,
+            )
+
+        with gr.Row():
+            submit_btn = gr.Button("Analyze", variant="primary")
+            clear_btn = gr.Button("Clear")
+
+        output_box = gr.Textbox(label="Response", lines=16)
+
+        # Wire events
+        submit_btn.click(fn=analyze_skin_condition, inputs=[image_input, question_input], outputs=output_box, queue=True)
+        clear_btn.click(fn=lambda: (None, ""), inputs=None, outputs=[image_input, question_input])
+
+        # Queue for concurrency control (ZeroGPU friendly)
+        demo.queue(concurrency_count=1, status_update_rate=1)
+
+        gr.Markdown(
+            "Tips: Ensure good lighting and focus. Avoid uploading personally identifying information."
         )
-        
-        # Model status
-        if model_loaded:
-            gr.HTML("<div style='text-align: center; color: green;'>✅ Model loaded successfully!</div>")
-        else:
-            gr.HTML("<div style='text-align: center; color: red;'>❌ Model loading failed. Please check the logs.</div>")
-    
+
     return demo
 
-@spaces.GPU
+
 def main():
-    """Main function with GPU decorator for Hugging Face Spaces"""
-    try:
-        # Create and launch the interface
-        demo = create_interface()
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            share=False,
-            show_error=True,
-            inbrowser=False,
-            quiet=False
-        )
-    except Exception as e:
-        logger.error(f"Error launching app: {e}")
-        raise
+    demo = create_interface()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True,
+        inbrowser=False,
+        quiet=False,
+        ssr_mode=False,  # disable SSR to avoid Node 20 requirement in Spaces container
+    )
+
 
 if __name__ == "__main__":
     main()
-