import gradio as gr import json, os, re, traceback, contextlib from typing import Any, List, Dict 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, ) -> 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() with amp_ctx: generated_ids = model.generate( **inputs, max_new_tokens=128, do_sample=False, ) 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 "" # --- 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): 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) Run inference try: coordinates_str = run_inference_localization(messages, resized_image, device, dtype) except Exception as e: traceback.print_exc() return f"Error during model inference: {e}", resized_image.copy().convert("RGB"), f"device: {device} | dtype: {dtype}" # 4) Parse coordinates and draw marker output_image_with_click = resized_image.copy().convert("RGB") match = re.search(r"Click\((\d+),\s*(\d+)\)", coordinates_str) if match: try: x = int(match.group(1)) y = int(match.group(2)) draw = ImageDraw.Draw(output_image_with_click) radius = max(5, min(resized_width // 100, resized_height // 100, 15)) bbox = (x - radius, y - radius, x + radius, y + radius) draw.ellipse(bbox, outline="red", width=max(2, radius // 4)) print(f"Predicted and drawn click at: ({x}, {y}) on resized image ({resized_width}x{resized_height})") except Exception as e: print(f"Error drawing on image: {e}") traceback.print_exc() else: print(f"Could not parse 'Click(x, y)' from model output: {coordinates_str}") return coordinates_str, output_image_with_click, f"device: {device} | dtype: {str(dtype).replace('torch.', '')}" # --- 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." ) submit_button = gr.Button("Localize Click", variant="primary") with gr.Column(scale=1): output_coords_component = gr.Textbox( label="Predicted Coordinates (Format: Click(x, y))", interactive=False ) output_image_component = gr.Image( type="pil", label="Image with Predicted Click Point", 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]], inputs=[input_image_component, instruction_component], 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], outputs=[output_coords_component, output_image_component, runtime_info] ) if __name__ == "__main__": # Do NOT pass 'concurrency_count' or ZeroGPU-specific launch args. demo.launch(debug=True)