import gradio as gr import json, os, re, traceback, contextlib, math, random from typing import Any, List, Dict, Optional, Tuple import spaces import torch from PIL import Image, ImageDraw import requests from transformers import AutoModelForImageTextToText, AutoProcessor from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize # --- Configuration --- MODEL_ID = "Hcompany/Holo1-3B" # ---------------- Device / DType helpers ---------------- def pick_device() -> str: """ On HF Spaces (ZeroGPU), CUDA is only available inside @spaces.GPU calls. We still honor FORCE_DEVICE for local testing. """ forced = os.getenv("FORCE_DEVICE", "").lower().strip() if forced in {"cpu", "cuda", "mps"}: return forced if torch.cuda.is_available(): return "cuda" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): return "mps" return "cpu" def pick_dtype(device: str) -> torch.dtype: if device == "cuda": major, _ = torch.cuda.get_device_capability() return torch.bfloat16 if major >= 8 else torch.float16 # Ampere+ -> bf16 if device == "mps": return torch.float16 return torch.float32 # CPU def move_to_device(batch, device: str): if isinstance(batch, dict): return {k: (v.to(device, non_blocking=True) if hasattr(v, "to") else v) for k, v in batch.items()} if hasattr(batch, "to"): return batch.to(device, non_blocking=True) return batch # --- Chat/template helpers --- def apply_chat_template_compat(processor, messages: List[Dict[str, Any]]) -> str: tok = getattr(processor, "tokenizer", None) if hasattr(processor, "apply_chat_template"): return processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) if tok is not None and hasattr(tok, "apply_chat_template"): return tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) texts = [] for m in messages: for c in m.get("content", []): if isinstance(c, dict) and c.get("type") == "text": texts.append(c.get("text", "")) return "\n".join(texts) def batch_decode_compat(processor, token_id_batches, **kw): tok = getattr(processor, "tokenizer", None) if tok is not None and hasattr(tok, "batch_decode"): return tok.batch_decode(token_id_batches, **kw) if hasattr(processor, "batch_decode"): return processor.batch_decode(token_id_batches, **kw) raise AttributeError("No batch_decode available on processor or tokenizer.") def get_image_proc_params(processor) -> Dict[str, int]: ip = getattr(processor, "image_processor", None) return { "patch_size": getattr(ip, "patch_size", 14), "merge_size": getattr(ip, "merge_size", 1), "min_pixels": getattr(ip, "min_pixels", 256 * 256), "max_pixels": getattr(ip, "max_pixels", 1280 * 1280), } def trim_generated(generated_ids, inputs): in_ids = getattr(inputs, "input_ids", None) if in_ids is None and isinstance(inputs, dict): in_ids = inputs.get("input_ids", None) if in_ids is None: return [out_ids for out_ids in generated_ids] return [out_ids[len(in_seq):] for in_seq, out_ids in zip(in_ids, generated_ids)] # --- Load model/processor ON CPU at import time (required for ZeroGPU) --- print(f"Loading model and processor for {MODEL_ID} on CPU startup (ZeroGPU safe)...") model = None processor = None model_loaded = False load_error_message = "" try: model = AutoModelForImageTextToText.from_pretrained( MODEL_ID, torch_dtype=torch.float32, # CPU-safe dtype at import trust_remote_code=True, ) processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) model.eval() model_loaded = True print("Model and processor loaded on CPU.") except Exception as e: load_error_message = ( f"Error loading model/processor: {e}\n" "This might be due to network/model ID/library versions.\n" "Check the full traceback in the logs." ) print(load_error_message) traceback.print_exc() # --- Prompt builder --- def get_localization_prompt(pil_image: Image.Image, instruction: str) -> List[dict]: guidelines: str = ( "Localize an element on the GUI image according to my instructions and " "output a click position as Click(x, y) with x num pixels from the left edge " "and y num pixels from the top edge." ) return [ { "role": "user", "content": [ {"type": "image", "image": pil_image}, {"type": "text", "text": f"{guidelines}\n{instruction}"} ], } ] # --- Inference core (device passed in; AMP used when suitable) --- @torch.inference_mode() def run_inference_localization( messages_for_template: List[dict[str, Any]], pil_image_for_processing: Image.Image, device: str, dtype: torch.dtype, do_sample: bool = False, temperature: float = 0.6, top_p: float = 0.9, max_new_tokens: int = 128, ) -> str: text_prompt = apply_chat_template_compat(processor, messages_for_template) inputs = processor( text=[text_prompt], images=[pil_image_for_processing], padding=True, return_tensors="pt", ) inputs = move_to_device(inputs, device) # AMP contexts if device == "cuda": amp_ctx = torch.autocast(device_type="cuda", dtype=dtype) elif device == "mps": amp_ctx = torch.autocast(device_type="mps", dtype=torch.float16) else: amp_ctx = contextlib.nullcontext() gen_kwargs = dict( max_new_tokens=max_new_tokens, do_sample=do_sample, temperature=temperature, top_p=top_p, ) with amp_ctx: generated_ids = model.generate(**inputs, **gen_kwargs) generated_ids_trimmed = trim_generated(generated_ids, inputs) decoded_output = batch_decode_compat( processor, generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False ) return decoded_output[0] if decoded_output else "" # ---------- Confidence helpers ---------- CLICK_RE = re.compile(r"Click\((\d+),\s*(\d+)\)") def parse_click(s: str) -> Optional[Tuple[int, int]]: m = CLICK_RE.search(s) if not m: return None try: return int(m.group(1)), int(m.group(2)) except Exception: return None @torch.inference_mode() def sample_clicks( messages: List[dict], img: Image.Image, device: str, dtype: torch.dtype, n_samples: int = 7, temperature: float = 0.6, top_p: float = 0.9, seed: Optional[int] = None, ) -> List[Optional[Tuple[int, int]]]: """ Run multiple stochastic decodes to estimate self-consistency. Returns a list of (x,y) or None (if parsing failed) for each sample. """ clicks: List[Optional[Tuple[int, int]]] = [] # If model respects torch random, set seed for reproducibility (optional) if seed is not None: torch.manual_seed(seed) random.seed(seed) for i in range(n_samples): # Vary seed slightly each iteration to avoid identical sampling patterns if seed is not None: torch.manual_seed(seed + i + 1) random.seed((seed + i + 1) & 0xFFFFFFFF) out = run_inference_localization( messages, img, device, dtype, do_sample=True, temperature=temperature, top_p=top_p ) clicks.append(parse_click(out)) return clicks def cluster_and_confidence( clicks: List[Optional[Tuple[int,int]]], img_w: int, img_h: int, ) -> Dict[str, Any]: """ Simple robust consensus: - Keep only valid points - Compute median point (x_med, y_med) - Compute distances to median - Inlier threshold = max(8 px, 2% of min(img_w, img_h)) - Confidence = (#inliers / #total_samples) * clamp(1 - (rms_inlier_dist / thr), 0, 1) Returns dict with consensus point, confidence, dispersion, and counts. """ valid = [xy for xy in clicks if xy is not None] total = len(clicks) if total == 0: return dict(ok=False, reason="no_samples") if not valid: return dict(ok=False, reason="no_valid_points", total=total) xs = sorted([x for x, _ in valid]) ys = sorted([y for _, y in valid]) mid = len(valid) // 2 if len(valid) % 2 == 1: x_med = xs[mid] y_med = ys[mid] else: x_med = (xs[mid - 1] + xs[mid]) // 2 y_med = (ys[mid - 1] + ys[mid]) // 2 thr = max(8.0, 0.02 * min(img_w, img_h)) # ~2% of smaller side, at least 8 px dists = [math.hypot(x - x_med, y - y_med) for (x, y) in valid] inliers = [(xy, d) for xy, d in zip(valid, dists) if d <= thr] outliers = [(xy, d) for xy, d in zip(valid, dists) if d > thr] inlier_count = len(inliers) # RMS of inlier distances (0 if perfect agreement) if inliers: rms = math.sqrt(sum(d*d for _, d in inliers) / len(inliers)) else: rms = float("inf") # Confidence: agreement ratio * sharpness factor if inliers: sharp = max(0.0, min(1.0, 1.0 - (rms / thr))) else: sharp = 0.0 confidence = (inlier_count / total) * sharp return dict( ok=True, x=x_med, y=y_med, confidence=confidence, total_samples=total, valid_samples=len(valid), inliers=inlier_count, outliers=len(outliers), sigma_px=rms if math.isfinite(rms) else None, inlier_threshold_px=thr, all_points=valid, inlier_points=[xy for xy,_ in inliers], outlier_points=[xy for xy,_ in outliers], ) def draw_samples( base_img: Image.Image, consensus_xy: Optional[Tuple[int,int]], inliers: List[Tuple[int,int]], outliers: List[Tuple[int,int]], ring_color: str = "red", ) -> Image.Image: """ Overlay all sampled points: green=inliers, red=outliers, plus a ring for consensus. """ img = base_img.copy().convert("RGB") draw = ImageDraw.Draw(img) w, h = img.size # Dot radius scales with image size r = max(3, min(w, h) // 200) # Draw inliers for (x, y) in inliers: draw.ellipse((x - r, y - r, x + r, y + r), fill="green", outline=None) # Draw outliers for (x, y) in outliers: draw.ellipse((x - r, y - r, x + r, y + r), fill="red", outline=None) # Consensus ring if consensus_xy is not None: cx, cy = consensus_xy ring_r = max(5, min(w, h) // 100, r * 3) draw.ellipse((cx - ring_r, cy - ring_r, cx + ring_r, cy + ring_r), outline=ring_color, width=max(2, ring_r // 4)) return img # --- Gradio processing function (ZeroGPU-visible) --- # Decorate the function Gradio calls so Spaces detects a GPU entry point. @spaces.GPU(duration=120) # keep GPU attached briefly between calls (seconds) def predict_click_location( input_pil_image: Image.Image, instruction: str, estimate_confidence: bool = True, num_samples: int = 7, temperature: float = 0.6, top_p: float = 0.9, seed: Optional[int] = None, ): if not model_loaded or not processor or not model: return f"Model not loaded. Error: {load_error_message}", None, "device: n/a | dtype: n/a" if not input_pil_image: return "No image provided. Please upload an image.", None, "device: n/a | dtype: n/a" if not instruction or instruction.strip() == "": return "No instruction provided. Please type an instruction.", input_pil_image.copy().convert("RGB"), "device: n/a | dtype: n/a" # Decide device/dtype *inside* the GPU-decorated call device = pick_device() dtype = pick_dtype(device) # Optional perf knobs for CUDA if device == "cuda": torch.backends.cuda.matmul.allow_tf32 = True torch.set_float32_matmul_precision("high") # If needed, move model now that GPU is available try: p = next(model.parameters()) cur_dev = p.device.type cur_dtype = p.dtype except StopIteration: cur_dev, cur_dtype = "cpu", torch.float32 if cur_dev != device or cur_dtype != dtype: model.to(device=device, dtype=dtype) model.eval() # 1) Resize according to image processor params (safe defaults if missing) try: ip = get_image_proc_params(processor) resized_height, resized_width = smart_resize( input_pil_image.height, input_pil_image.width, factor=ip["patch_size"] * ip["merge_size"], min_pixels=ip["min_pixels"], max_pixels=ip["max_pixels"], ) resized_image = input_pil_image.resize( size=(resized_width, resized_height), resample=Image.Resampling.LANCZOS ) except Exception as e: traceback.print_exc() return f"Error resizing image: {e}", input_pil_image.copy().convert("RGB"), f"device: {device} | dtype: {dtype}" # 2) Build messages with image + instruction messages = get_localization_prompt(resized_image, instruction) # 3) Inference and (optionally) confidence estimation try: if estimate_confidence and num_samples >= 3: # Monte-Carlo sampling clicks = sample_clicks( messages, resized_image, device, dtype, n_samples=int(num_samples), temperature=float(temperature), top_p=float(top_p), seed=seed ) summary = cluster_and_confidence(clicks, resized_image.width, resized_image.height) if not summary.get("ok", False): # Fallback: deterministic decode coord_str = run_inference_localization(messages, resized_image, device, dtype, do_sample=False) out_img = resized_image.copy().convert("RGB") match = CLICK_RE.search(coord_str or "") if match: x, y = int(match.group(1)), int(match.group(2)) out_img = draw_samples(out_img, (x, y), [], []) coords_text = f"{coord_str} | confidence=0.00 (fallback)" return coords_text, out_img, f"device: {device} | dtype: {str(dtype).replace('torch.', '')}" # Build final string + visualization x, y = int(summary["x"]), int(summary["y"]) conf = summary["confidence"] inliers = summary["inlier_points"] outliers = summary["outlier_points"] sigma = summary["sigma_px"] thr = summary["inlier_threshold_px"] total = summary["total_samples"] valid = summary["valid_samples"] # Compose output string in the same canonical format plus diagnostics coord_str = f"Click({x}, {y})" diag = ( f"confidence={conf:.2f} | samples(valid/total)={valid}/{total} | " f"inliers={len(inliers)} | σ={sigma:.1f}px | thr={thr:.1f}px | " f"T={temperature:.2f}, p={top_p:.2f}" ) # Draw: all samples + consensus ring out_img = draw_samples(resized_image, (x, y), inliers, outliers) return f"{coord_str} | {diag}", out_img, f"device: {device} | dtype: {str(dtype).replace('torch.', '')}" else: # Fast deterministic single pass (no confidence) coord_str = run_inference_localization(messages, resized_image, device, dtype, do_sample=False) out_img = resized_image.copy().convert("RGB") match = CLICK_RE.search(coord_str or "") if match: x = int(match.group(1)) y = int(match.group(2)) # draw a simple ring around the predicted point out_img = draw_samples(out_img, (x, y), [], []) else: print(f"Could not parse 'Click(x, y)' from model output: {coord_str}") return coord_str, out_img, f"device: {device} | dtype: {str(dtype).replace('torch.', '')}" except Exception as e: traceback.print_exc() return f"Error during model inference: {e}", resized_image.copy().convert("RGB"), f"device: {device} | dtype: {dtype}" # --- Load Example Data --- example_image = None example_instruction = "Enter the server address readyforquantum.com to check its security" try: example_image_url = "https://readyforquantum.com/img/screentest.jpg" example_image = Image.open(requests.get(example_image_url, stream=True).raw) except Exception as e: print(f"Could not load example image from URL: {e}") traceback.print_exc() try: example_image = Image.new("RGB", (200, 150), color="lightgray") draw = ImageDraw.Draw(example_image) draw.text((10, 10), "Example image\nfailed to load", fill="black") except Exception: pass # --- Gradio UI --- title = "Holo1-3B: Holo1 Localization Demo (ZeroGPU-ready)" article = f"""

Model: {MODEL_ID} by HCompany | Paper: HCompany Tech Report | Blog: Surfer-H Blog Post
GPU (if available) is requested only during inference via @spaces.GPU.

""" if not model_loaded: with gr.Blocks() as demo: gr.Markdown(f"#
⚠️ Error: Model Failed to Load ⚠️
") gr.Markdown(f"
{load_error_message}
") gr.Markdown("
See logs for the full traceback.
") else: with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown(f"

{title}

") gr.Markdown(article) with gr.Row(): with gr.Column(scale=1): input_image_component = gr.Image(type="pil", label="Input UI Image", height=400) instruction_component = gr.Textbox( label="Instruction", placeholder="e.g., Click the 'Login' button", info="Type the action you want the model to localize on the image." ) estimate_conf = gr.Checkbox(value=True, label="Estimate confidence (slower)") num_samples_slider = gr.Slider(3, 15, value=7, step=1, label="Samples (for confidence)") temperature_slider = gr.Slider(0.2, 1.2, value=0.6, step=0.05, label="Temperature") top_p_slider = gr.Slider(0.5, 0.99, value=0.9, step=0.01, label="Top-p") seed_box = gr.Number(value=None, precision=0, label="Seed (optional, for reproducibility)") submit_button = gr.Button("Localize Click", variant="primary") with gr.Column(scale=1): output_coords_component = gr.Textbox( label="Predicted Coordinates + Confidence", interactive=False ) output_image_component = gr.Image( type="pil", label="Image with Samples (green=inliers, red=outliers) and Final Ring", height=400, interactive=False ) runtime_info = gr.Textbox( label="Runtime Info", value="device: n/a | dtype: n/a", interactive=False ) if example_image: gr.Examples( examples=[[example_image, example_instruction, True, 7, 0.6, 0.9, None]], inputs=[ input_image_component, instruction_component, estimate_conf, num_samples_slider, temperature_slider, top_p_slider, seed_box, ], outputs=[output_coords_component, output_image_component, runtime_info], fn=predict_click_location, cache_examples="lazy", ) submit_button.click( fn=predict_click_location, inputs=[ input_image_component, instruction_component, estimate_conf, num_samples_slider, temperature_slider, top_p_slider, seed_box, ], outputs=[output_coords_component, output_image_component, runtime_info] ) if __name__ == "__main__": demo.launch(debug=True)