Update app.py

app.py

@@ -1,8 +1,9 @@
 # app.py
 # Dermatology-AI-Assistant — Hugging Face Space (ZeroGPU-ready)
-# - Uses qwen-vl-utils for vision inputs
+# - First tries your fine-tuned model
+# - If Qwen raises token/feature mismatch, falls back to official base model
 # - Acquires ZeroGPU only during inference
-# - Handles Qwen2-VL token/feature mismatch with a safe fallback retry
+# - Uses qwen-vl-utils.process_vision_info
 
 import os
 import logging
@@ -15,42 +16,47 @@ from PIL import Image
 from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
 from qwen_vl_utils import process_vision_info
 
-# ---------------------------
-# Logging
-# ---------------------------
 logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
 logger = logging.getLogger(__name__)
 
 # ---------------------------
 # Config
 # ---------------------------
-MODEL_ID = os.environ.get("MODEL_ID", "ColdSlim/Dermatology-Qwen2.5-VL-3B")
+FT_MODEL_ID = os.environ.get("MODEL_ID", "ColdSlim/Dermatology-Qwen2.5-VL-3B")
+BASE_MODEL_ID = os.environ.get("FALLBACK_BASE_MODEL_ID", "Qwen/Qwen2.5-VL-3B-Instruct")
+
 GEN_KW = dict(
     max_new_tokens=512,
     do_sample=True,
     temperature=0.7,
     top_p=0.9,
 )
+
 ZGPU_DURATION = int(os.environ.get("ZGPU_DURATION", "180"))
 
-logger.info(f"Loading processor from: {MODEL_ID}")
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+# Preload only the fine-tuned processor on CPU; we may swap to base processor in the fallback
+logger.info(f"Loading processor from: {FT_MODEL_ID}")
+ft_processor = AutoProcessor.from_pretrained(FT_MODEL_ID, trust_remote_code=True)
+logger.info("Processor loaded.")
 
-# (Optional) Tame resolution to reduce tiling variance; adjust if you like.
-if hasattr(processor, "image_processor"):
-    try:
-        # Keep images within a predictable pixel band so placeholder count is stable.
-        processor.image_processor.max_pixels = int(os.environ.get("QWEN_MAX_PIXELS", "1500000"))  # ~1.5MP
-        processor.image_processor.min_pixels = int(os.environ.get("QWEN_MIN_PIXELS", "262144"))   # 512x512
-    except Exception:
-        pass
+# Optional: stabilize tiling by constraining pixel range (helps placeholder consistency)
+def _tune_image_processor(proc):
+    if hasattr(proc, "image_processor"):
+        try:
+            proc.image_processor.max_pixels = int(os.environ.get("QWEN_MAX_PIXELS", "1500000"))  # ~1.5MP
+            proc.image_processor.min_pixels = int(os.environ.get("QWEN_MIN_PIXELS", "262144"))   # 512x512
+        except Exception:
+            pass
 
-logger.info("Processor loaded.")
+_tune_image_processor(ft_processor)
 
 # ---------------------------
 # Helpers
 # ---------------------------
 def _messages(image: Image.Image, question: str):
+    # ensure RGB to avoid mode surprises
+    if image.mode != "RGB":
+        image = image.convert("RGB")
     return [
         {
             "role": "user",
@@ -61,47 +67,28 @@ def _messages(image: Image.Image, question: str):
         }
     ]
 
-def build_inputs(image: Image.Image, question: str, *, disable_splitting: bool = False):
+def build_inputs(processor: AutoProcessor, image: Image.Image, question: str):
     """
-    Build Qwen-style multimodal chat inputs.
-    When disable_splitting is True, we hint the image processor to avoid tiling,
-    which can fix token/feature mismatches for some edge cases.
+    Build Qwen-style multimodal inputs (no padding, batch size 1).
     """
     messages = _messages(image, question)
-
-    # Apply chat template (inserts <image> placeholders automatically)
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-    # Prepare vision inputs
     image_inputs, video_inputs = process_vision_info(messages)
-
-    # Optionally force-disable splitting (fallback path)
-    if disable_splitting and hasattr(processor, "image_processor"):
-        ip = processor.image_processor
-        # Cache old setting to not mutate global defaults permanently
-        prev = getattr(ip, "do_image_splitting", None)
-        try:
-            if hasattr(ip, "do_image_splitting"):
-                ip.do_image_splitting = False
-            inputs = processor(
-                text=[text],
-                images=image_inputs,
-                videos=video_inputs,
-                return_tensors="pt",   # <- no padding for single-sample path
-            )
-        finally:
-            if prev is not None:
-                ip.do_image_splitting = prev
-    else:
-        inputs = processor(
-            text=[text],
-            images=image_inputs,
-            videos=video_inputs,
-            return_tensors="pt",       # <- no padding to avoid mask quirks
-        )
-
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        return_tensors="pt",  # no padding for single sample
+    )
     return inputs
 
+def _pad_token_id(processor, model):
+    # Prefer tokenizer.eos if present; else model config; else 0
+    tid = getattr(getattr(processor, "tokenizer", None), "eos_token_id", None)
+    if tid is not None:
+        return tid
+    return getattr(getattr(model, "config", None), "eos_token_id", 0)
+
 def format_derm_disclaimer(ans: str) -> str:
     tail = (
         "\n\n---\n"
@@ -110,69 +97,81 @@ def format_derm_disclaimer(ans: str) -> str:
     )
     return ans + tail
 
+def _generate_text(model, processor, inputs: dict) -> str:
+    # move to CUDA
+    inputs = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+    with torch.no_grad():
+        out_ids = model.generate(
+            **inputs,
+            **GEN_KW,
+            pad_token_id=_pad_token_id(processor, model),
+        )
+    trimmed = [o[len(i):] for i, o in zip(inputs["input_ids"], out_ids)]
+    text = processor.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    return text
+
 # ---------------------------
 # Inference (ZeroGPU)
 # ---------------------------
 @spaces.GPU(duration=ZGPU_DURATION)
 def analyze_skin_condition(image: Optional[Image.Image], question: str) -> str:
     """
-    Runs inside a ZeroGPU reservation window.
-    Loads model on GPU, generates, frees VRAM.
-    Includes a fallback retry if Qwen raises a token/feature mismatch.
+    Try fine-tuned model first; on token/feature mismatch, fall back to base model+processor.
     """
     if image is None:
         return "❌ Please upload an image first."
 
+    model = None
     try:
-        logger.info(f"Loading model on GPU: {MODEL_ID}")
+        # ------- Attempt 1: Fine-tuned model -------
+        logger.info(f"Loading fine-tuned model on GPU: {FT_MODEL_ID}")
         model = Qwen2VLForConditionalGeneration.from_pretrained(
-            MODEL_ID,
+            FT_MODEL_ID,
             torch_dtype=torch.float16,
             device_map="cuda",
             trust_remote_code=True,
             low_cpu_mem_usage=True,
-            ignore_mismatched_sizes=True,  # keep until weights align perfectly
+            ignore_mismatched_sizes=True,  # your FT ckpt logs suggest some vision head diffs
         )
-        logger.info("Model loaded successfully!")
-
-        def _run_infer(disable_splitting: bool = False) -> str:
-            inputs = build_inputs(image, question, disable_splitting=disable_splitting)
-            # Move tensors to CUDA
-            inputs = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
-            with torch.no_grad():
-                out_ids = model.generate(
-                    **inputs,
-                    **GEN_KW,
-                    pad_token_id=processor.tokenizer.eos_token_id,
-                )
-            # Strip prompt tokens before decoding
-            trimmed = [o[len(i):] for i, o in zip(inputs["input_ids"], out_ids)]
-            text = processor.batch_decode(
-                trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-            )[0]
-            return text
-
-        # First attempt: normal path
+        logger.info("Fine-tuned model loaded.")
+        inputs = build_inputs(ft_processor, image, question)
         try:
-            text = _run_infer(disable_splitting=False)
+            text = _generate_text(model, ft_processor, inputs)
+            return format_derm_disclaimer(text)
         except ValueError as ve:
             msg = str(ve)
-            # Known Qwen2-VL edge case: token/feature mismatch — retry with splitting disabled
             if "Image features and image tokens do not match" in msg:
-                logger.warning("Token/feature mismatch detected — retrying with image splitting disabled.")
-                text = _run_infer(disable_splitting=True)
+                logger.warning("Token/feature mismatch on fine-tuned model — falling back to base model.")
             else:
                 raise
 
-        # Free VRAM
+        # ------- Attempt 2: Base model & its processor -------
+        # Free FT model first
         del model
         torch.cuda.empty_cache()
 
+        logger.info(f"Loading BASE model on GPU: {BASE_MODEL_ID}")
+        base_processor = AutoProcessor.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
+        _tune_image_processor(base_processor)
+        model = Qwen2VLForConditionalGeneration.from_pretrained(
+            BASE_MODEL_ID,
+            torch_dtype=torch.float16,
+            device_map="cuda",
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        logger.info("Base model loaded.")
+        base_inputs = build_inputs(base_processor, image, question)
+        text = _generate_text(model, base_processor, base_inputs)
         return format_derm_disclaimer(text)
 
     except Exception as e:
         logger.exception("Error during inference")
         return f"❌ Error analyzing image: {e}"
+    finally:
+        if model is not None:
+            del model
+        torch.cuda.empty_cache()
 
 # ---------------------------
 # UI
@@ -201,9 +200,7 @@ def create_interface() -> gr.Blocks:
         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])
 
-        # Gradio 4.44.1: simple queue() call (no kwargs)
         demo.queue()
-
         gr.Markdown("Tips: Ensure good lighting and focus. Avoid uploading personally identifying information.")
     return demo
 
@@ -216,7 +213,7 @@ def main():
         show_error=True,
         inbrowser=False,
         quiet=False,
-        ssr_mode=False,  # avoid Node 20 requirement in container
+        ssr_mode=False,
     )
 
 if __name__ == "__main__":