diff --git a/modules/interrogate/moondream3.py b/modules/interrogate/moondream3.py new file mode 100644 index 000000000..ad9214fa4 --- /dev/null +++ b/modules/interrogate/moondream3.py @@ -0,0 +1,443 @@ +# Moondream 3 Preview VLM Implementation +# Source: https://huggingface.co/moondream/moondream3-preview +# Model: 9.3GB, gated (requires HuggingFace authentication) +# Architecture: Mixture-of-Experts (9B total params, 2B active) +import os +import re +import torch +import transformers +from PIL import Image +from modules import shared, devices, sd_models + + +# Debug logging - function-based to avoid circular import +debug_enabled = os.environ.get('SD_VQA_DEBUG', None) is not None + +def debug(*args, **kwargs): + if debug_enabled: + shared.log.trace(*args, **kwargs) + + +# Global state +moondream3_model = None +loaded = None +image_cache = {} # Cache encoded images for reuse + + +def get_settings(): + """ + Build settings dict for Moondream 3 API from global VQA options. + Moondream 3 accepts: temperature, top_p, max_tokens + """ + settings = {} + if shared.opts.interrogate_vlm_max_length > 0: + settings['max_tokens'] = shared.opts.interrogate_vlm_max_length + if shared.opts.interrogate_vlm_temperature > 0: + settings['temperature'] = shared.opts.interrogate_vlm_temperature + if shared.opts.interrogate_vlm_top_p > 0: + settings['top_p'] = shared.opts.interrogate_vlm_top_p + return settings if settings else None + + +def load_model(repo: str): + """Load and compile Moondream 3 model.""" + global moondream3_model, loaded # pylint: disable=global-statement + + if moondream3_model is None or loaded != repo: + shared.log.debug(f'Interrogate load: vlm="{repo}"') + moondream3_model = None + + moondream3_model = transformers.AutoModelForCausalLM.from_pretrained( + repo, + trust_remote_code=True, + torch_dtype=devices.dtype, + cache_dir=shared.opts.hfcache_dir, + ) + + moondream3_model.eval() + if 'LLM' in shared.opts.cuda_compile: + debug('VQA interrogate: handler=moondream3 compiling model for fast decoding') + moondream3_model.compile() # Critical for fast decoding per moondream3 docs + loaded = repo + devices.torch_gc() + + # Move model to active device + sd_models.move_model(moondream3_model, devices.device) + return moondream3_model + + +def encode_image(image: Image.Image, cache_key: str = None): + """ + Encode image for reuse across multiple queries. + + Args: + image: PIL Image + cache_key: Optional cache key for storing encoded image + + Returns: + Encoded image tensor + """ + if cache_key and cache_key in image_cache: + debug(f'VQA interrogate: handler=moondream3 using cached encoding for cache_key="{cache_key}"') + return image_cache[cache_key] + + model = load_model(loaded) + + with devices.inference_context(): + encoded = model.encode_image(image) + + if cache_key: + image_cache[cache_key] = encoded + debug(f'VQA interrogate: handler=moondream3 cached encoding cache_key="{cache_key}" cache_size={len(image_cache)}') + + return encoded + + +def query(image: Image.Image, question: str, repo: str, stream: bool = False, + temperature: float = None, top_p: float = None, max_tokens: int = None, + use_cache: bool = False, reasoning: bool = True): + """ + Visual question answering with optional streaming. + + Args: + image: PIL Image + question: Question about the image + repo: Model repository + stream: Enable streaming output (generator) + temperature: Sampling temperature (overrides global setting) + top_p: Nucleus sampling parameter (overrides global setting) + max_tokens: Maximum tokens to generate (overrides global setting) + use_cache: Use cached image encoding if available + + Returns: + Answer dict or string (or generator if stream=True) + """ + model = load_model(repo) + + # Build settings - per-call parameters override global settings + settings = get_settings() or {} + if temperature is not None: + settings['temperature'] = temperature + if top_p is not None: + settings['top_p'] = top_p + if max_tokens is not None: + settings['max_tokens'] = max_tokens + + debug(f'VQA interrogate: handler=moondream3 method=query question="{question}" stream={stream} settings={settings}') + + # Use cached encoding if requested + if use_cache: + cache_key = f"{id(image)}_{question}" + image_input = encode_image(image, cache_key) + else: + image_input = image + + with devices.inference_context(): + response = model.query( + image=image_input, + question=question, + stream=stream, + settings=settings if settings else None, + reasoning=reasoning + ) + + # Log response structure (for non-streaming) + if not stream: + if isinstance(response, dict): + debug(f'VQA interrogate: handler=moondream3 response_type=dict keys={list(response.keys())}') + if 'reasoning' in response: + reasoning_text = response['reasoning'].get('text', '')[:100] + '...' if len(response['reasoning'].get('text', '')) > 100 else response['reasoning'].get('text', '') + debug(f'VQA interrogate: handler=moondream3 reasoning="{reasoning_text}"') + if 'answer' in response: + debug(f'VQA interrogate: handler=moondream3 answer="{response["answer"]}"') + + return response + + +def caption(image: Image.Image, repo: str, length: str = 'normal', stream: bool = False, + temperature: float = None, top_p: float = None, max_tokens: int = None): + """ + Generate image captions at different lengths. + + Args: + image: PIL Image + repo: Model repository + length: Caption length - 'short', 'normal', or 'long' + stream: Enable streaming output (generator) + temperature: Sampling temperature (overrides global setting) + top_p: Nucleus sampling parameter (overrides global setting) + max_tokens: Maximum tokens to generate (overrides global setting) + + Returns: + Caption dict or string (or generator if stream=True) + """ + model = load_model(repo) + + # Build settings - per-call parameters override global settings + settings = get_settings() or {} + if temperature is not None: + settings['temperature'] = temperature + if top_p is not None: + settings['top_p'] = top_p + if max_tokens is not None: + settings['max_tokens'] = max_tokens + + debug(f'VQA interrogate: handler=moondream3 method=caption length={length} stream={stream} settings={settings}') + + with devices.inference_context(): + response = model.caption( + image, + length=length, + stream=stream, + settings=settings if settings else None + ) + + # Log response structure (for non-streaming) + if not stream and isinstance(response, dict): + debug(f'VQA interrogate: handler=moondream3 response_type=dict keys={list(response.keys())}') + + return response + + +def point(image: Image.Image, object_name: str, repo: str): + """ + Identify coordinates of all instances of a specific object in the image. + + Args: + image: PIL Image + object_name: Name of object to locate + repo: Model repository + + Returns: + List of (x, y) tuples with coordinates normalized to 0-1 range, or None if not found + Example: [(0.733, 0.442), (0.5, 0.6)] for 2 instances + """ + model = load_model(repo) + + debug(f'VQA interrogate: handler=moondream3 method=point object_name="{object_name}"') + + with devices.inference_context(): + result = model.point(image, object_name) + + # Debug: Log the actual result to understand the format + debug(f'VQA interrogate: handler=moondream3 point_raw_result="{result}" type={type(result)}') + if isinstance(result, dict): + debug(f'VQA interrogate: handler=moondream3 point_raw_result_keys={list(result.keys())}') + + # Parse and validate coordinates + # Handle dict format: {'points': [{'x': 0.733, 'y': 0.442}, {'x': 0.5, 'y': 0.6}, ...]} + if isinstance(result, dict) and 'points' in result: + points_list = result['points'] + if points_list and len(points_list) > 0: + coordinates = [] + for point_data in points_list: # Iterate ALL points + if 'x' in point_data and 'y' in point_data: + x = max(0.0, min(1.0, float(point_data['x']))) + y = max(0.0, min(1.0, float(point_data['y']))) + coordinates.append((x, y)) + if coordinates: + debug(f'VQA interrogate: handler=moondream3 point_result={len(coordinates)} points found') + return coordinates + # Fallback: try simple list/tuple format [x, y] (for compatibility) + elif isinstance(result, (list, tuple)) and len(result) == 2: + x, y = result + x = max(0.0, min(1.0, float(x))) + y = max(0.0, min(1.0, float(y))) + debug('VQA interrogate: handler=moondream3 point_result=1 point found') + return [(x, y)] # Return as list for consistency + + debug('VQA interrogate: handler=moondream3 point_result=not found') + return None + + +def detect(image: Image.Image, object_name: str, repo: str, max_objects: int = 10): + """ + Detect all instances of a specific object with bounding boxes. + + Args: + image: PIL Image + object_name: Name of object to detect + repo: Model repository + max_objects: Maximum number of objects to return + + Returns: + List of detection dicts with keys: + - 'bbox': [x1, y1, x2, y2] normalized to 0-1 + - 'label': Object label + - 'confidence': Detection confidence (0-1) + Returns empty list if no objects found. + """ + model = load_model(repo) + + debug(f'VQA interrogate: handler=moondream3 method=detect object_name="{object_name}" max_objects={max_objects}') + + with devices.inference_context(): + result = model.detect(image, object_name) + + # Debug: Log the actual result to understand the format + debug(f'VQA interrogate: handler=moondream3 detect_raw_result="{result}" type={type(result)}') + if isinstance(result, dict): + debug(f'VQA interrogate: handler=moondream3 detect_raw_result_keys={list(result.keys())}') + + # Parse detections + # Expected format: {'objects': [{'x_min': 0.1, 'y_min': 0.2, 'x_max': 0.5, 'y_max': 0.8}, ...]} + detections = [] + + if isinstance(result, dict) and 'objects' in result: + objects = result['objects'][:max_objects] # Limit to max_objects + for i, obj in enumerate(objects): + if all(k in obj for k in ['x_min', 'y_min', 'x_max', 'y_max']): + bbox = [ + max(0.0, min(1.0, float(obj['x_min']))), + max(0.0, min(1.0, float(obj['y_min']))), + max(0.0, min(1.0, float(obj['x_max']))), + max(0.0, min(1.0, float(obj['y_max']))) + ] + detections.append({ + 'bbox': bbox, + 'label': object_name, + 'confidence': obj.get('confidence', 1.0) # Default confidence if not provided + }) + + debug(f'VQA interrogate: handler=moondream3 detect_result={len(detections)} objects found') + return detections + + +def predict(question: str, image: Image.Image, repo: str, model_name: str = None, thinking_mode: bool = False, + mode: str = None, stream: bool = False, use_cache: bool = False, **kwargs): + """ + Main entry point for Moondream 3 VQA - auto-detects mode from question. + + Args: + question: The question/prompt (e.g., "caption", "where is the cat?", "describe this") + image: PIL Image + repo: Model repository + model_name: Display name for logging + thinking_mode: Enable reasoning mode for query + mode: Force specific mode ('query', 'caption', 'caption_short', 'caption_long', 'point', 'detect') + stream: Enable streaming output (for query/caption) + use_cache: Use cached image encoding (for query) + **kwargs: Additional parameters (max_objects for detect, etc.) + + Returns: + Response string or tuple (text, annotated_image) for detect/point modes + (or generator if stream=True for query/caption modes) + """ + debug(f'VQA interrogate: handler=moondream3 model_name="{model_name}" repo="{repo}" question="{question}" image_size={image.size if image else None} mode={mode} stream={stream}') + + # Clean question + question = question.replace('<', '').replace('>', '').replace('_', ' ') if question else '' + + # Auto-detect mode from question if not specified + if mode is None: + question_lower = question.lower() + + # Caption detection + if question in ['CAPTION', 'caption'] or 'caption' in question_lower: + if 'more detailed' in question_lower or 'very long' in question_lower: + mode = 'caption_long' + elif 'detailed' in question_lower or 'long' in question_lower: + mode = 'caption_normal' + elif 'short' in question_lower or 'brief' in question_lower: + mode = 'caption_short' + else: + # Default caption mode (matches vqa.py legacy behavior) + if question == 'CAPTION': + mode = 'caption_short' + elif question == 'DETAILED CAPTION': + mode = 'caption_normal' + elif question == 'MORE DETAILED CAPTION': + mode = 'caption_long' + else: + mode = 'caption_normal' + + # Point detection + elif 'where is' in question_lower or 'locate' in question_lower or 'find' in question_lower or 'point' in question_lower: + mode = 'point' + + # Object detection + elif 'detect' in question_lower or 'bounding box' in question_lower or 'bbox' in question_lower: + mode = 'detect' + + # Default to query + else: + mode = 'query' + + debug(f'VQA interrogate: handler=moondream3 mode_selected={mode}') + + # Dispatch to appropriate method + try: + if mode == 'caption_short': + response = caption(image, repo, length='short', stream=stream) + elif mode == 'caption_long': + response = caption(image, repo, length='long', stream=stream) + elif mode in ['caption', 'caption_normal']: + response = caption(image, repo, length='normal', stream=stream) + elif mode == 'point': + # Extract object name from question - case insensitive, preserve object names + object_name = question + # Remove trigger phrases (case-insensitive) + for phrase in ['point at', 'where is', 'locate', 'find']: + object_name = re.sub(rf'\b{phrase}\b', '', object_name, flags=re.IGNORECASE) + # Remove punctuation and extra whitespace + object_name = re.sub(r'[?.!,]', '', object_name).strip() + # Remove leading "the" only + object_name = re.sub(r'^\s*the\s+', '', object_name, flags=re.IGNORECASE) + debug(f'VQA interrogate: handler=moondream3 point_extracted_object="{object_name}"') + result = point(image, object_name, repo) + if result: + # Handle multiple instances - return text and points for drawing + if len(result) == 1: + text = f"Found at coordinates: ({result[0][0]:.3f}, {result[0][1]:.3f})" + else: + # Multiple instances found - format with count + lines = [f"Found {len(result)} instances:"] + for i, (x, y) in enumerate(result, 1): + lines.append(f" {i}. ({x:.3f}, {y:.3f})") + text = '\n'.join(lines) + return (text, {'points': result}) # Return text and points data + return ("Object not found", None) + elif mode == 'detect': + # Extract object name from question - case insensitive + object_name = question + # Remove trigger phrases (case-insensitive) + for phrase in ['detect', 'find all', 'bounding box', 'bbox', 'find']: + object_name = re.sub(rf'\b{phrase}\b', '', object_name, flags=re.IGNORECASE) + # Remove punctuation and extra whitespace + object_name = re.sub(r'[?.!,]', '', object_name).strip() + # Remove leading "the" only + object_name = re.sub(r'^\s*the\s+', '', object_name, flags=re.IGNORECASE) + # Remove "and" and get first object (model detects one type at a time) + if ' and ' in object_name.lower(): + object_name = re.split(r'\s+and\s+', object_name, flags=re.IGNORECASE)[0].strip() + debug(f'VQA interrogate: handler=moondream3 detect_extracted_object="{object_name}"') + + results = detect(image, object_name, repo, max_objects=kwargs.get('max_objects', 10)) + # Format as string for display and return detections for drawing + if results: + lines = [f"{det['label']}: [{det['bbox'][0]:.3f}, {det['bbox'][1]:.3f}, {det['bbox'][2]:.3f}, {det['bbox'][3]:.3f}] (confidence: {det['confidence']:.2f})" + for det in results] + text = '\n'.join(lines) + return (text, {'detections': results}) # Return text and detection data + return ("No objects detected", None) + else: # mode == 'query' + if len(question) < 2: + question = "Describe this image." + response = query(image, question, repo, stream=stream, use_cache=use_cache, reasoning=thinking_mode) + + debug(f'VQA interrogate: handler=moondream3 response_before_clean="{response}"') + return response + + except Exception as e: + from modules import errors + errors.display(e, 'Moondream3') + return f"Error: {str(e)}" + + +def clear_cache(): + """Clear image encoding cache.""" + global image_cache # pylint: disable=global-statement + cache_size = len(image_cache) + image_cache.clear() + debug(f'VQA interrogate: handler=moondream3 cleared image cache cache_size_was={cache_size}') + shared.log.debug(f'Moondream3: Cleared image cache ({cache_size} entries)')