Update app.py and requirements following Mungert's proposed setup

app.py

@@ -1,24 +1,89 @@
-import subprocess
-#subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
 import gradio as gr
-import json
-import os
-from typing import Any, List
-import spaces
+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
-import torch
-import re 
 
 # --- Configuration ---
-MODEL_ID = "Hcompany/Holo1-7B"
+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.")
 
-# --- Model and Processor Loading (Load once) ---
-print(f"Loading model and processor for {MODEL_ID}...")
+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
@@ -27,200 +92,185 @@ load_error_message = ""
 try:
     model = AutoModelForImageTextToText.from_pretrained(
         MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        attn_implementation="flash_attention_2",
-        trust_remote_code=True 
-    ).to("cuda")
+        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 successfully.")
+    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 issues, an incorrect model ID, or missing dependencies (like flash_attention_2 if enabled by default in some config).\n" \
-                         "Ensure you have a stable internet connection and the necessary libraries installed."
+    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()
 
-# --- Helper functions from the model card (or adapted) ---
-
-def get_localization_prompt(pil_image: Image.Image, instruction: str) -> List[dict[str, Any]]:
-    """
-    Prepares the prompt structure for the Holo1 model for localization tasks.
-    The `pil_image` argument here is primarily for semantic completeness in the prompt structure,
-    as the actual image tensor is handled by the processor later.
-    """
-    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."
-    
+# --- 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}"},
+                {"type": "image", "image": pil_image},
+                {"type": "text", "text": f"{guidelines}\n{instruction}"}
             ],
         }
     ]
 
-@spaces.GPU(duration=20)
+# --- 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
+    messages_for_template: List[dict[str, Any]],
+    pil_image_for_processing: Image.Image,
+    device: str,
+    dtype: torch.dtype,
 ) -> str:
-    model.to("cuda")
-    torch.cuda.set_device(0)
-    """
-    Runs inference using the Holo1 model.
-    - messages_for_template: The prompt structure, potentially including the PIL image object 
-                             (which apply_chat_template converts to an image tag).
-    - pil_image_for_processing: The actual PIL image to be processed into tensors.
-    """
-    # 1. Apply chat template to messages. This will create the text part of the prompt,
-    #    including image tags if the image was part of `messages_for_template`.
-    text_prompt = processor.apply_chat_template(
-        messages_for_template,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+    text_prompt = apply_chat_template_compat(processor, messages_for_template)
 
-    # 2. Process text and image together to get model inputs
     inputs = processor(
         text=[text_prompt],
-        images=[pil_image_for_processing], # Provide the actual image data here
+        images=[pil_image_for_processing],
         padding=True,
         return_tensors="pt",
     )
-    inputs = inputs.to(model.device)
-
-    # 3. Generate response
-    # Using do_sample=False for more deterministic output, as in the model card's structured output example
-    generated_ids = model.generate(**inputs, max_new_tokens=128, do_sample=False) 
-    
-    # 4. Trim input_ids from generated_ids to get only the generated part
-    generated_ids_trimmed = [
-        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    
-    # 5. Decode the generated tokens
-    decoded_output = processor.batch_decode(
-        generated_ids_trimmed, 
-        skip_special_tokens=True, 
+    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 ---
+# --- 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
+        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
+        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")
+        return "No instruction provided. Please type an instruction.", input_pil_image.copy().convert("RGB"), "device: n/a | dtype: n/a"
 
-    # 1. Prepare image: Resize according to model's image processor's expected properties
-    #    This ensures predicted coordinates match the (resized) image dimensions.
-    image_proc_config = processor.image_processor
+    # 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=image_proc_config.patch_size * image_proc_config.merge_size,
-            min_pixels=image_proc_config.min_pixels,
-            max_pixels=image_proc_config.max_pixels,
+            factor=ip["patch_size"] * ip["merge_size"],
+            min_pixels=ip["min_pixels"],
+            max_pixels=ip["max_pixels"],
         )
-        # Using LANCZOS for resampling as it's generally good for downscaling.
-        # The model card used `resample=None`, which might imply nearest or default.
-        # For visual quality in the demo, LANCZOS is reasonable.
         resized_image = input_pil_image.resize(
-            size=(resized_width, resized_height), 
-            resample=Image.Resampling.LANCZOS # type: ignore
+            size=(resized_width, resized_height),
+            resample=Image.Resampling.LANCZOS
         )
     except Exception as e:
-        print(f"Error resizing image: {e}")
-        return f"Error resizing image: {e}", input_pil_image.copy().convert("RGB")
+        traceback.print_exc()
+        return f"Error resizing image: {e}", input_pil_image.copy().convert("RGB"), f"device: {device} | dtype: {dtype}"
 
-    # 2. Create the prompt using the resized image (for correct image tagging context) and instruction
+    # 2) Build messages with image + instruction
     messages = get_localization_prompt(resized_image, instruction)
 
-    # 3. Run inference
-    #    Pass `messages` (which includes the image object for template processing)
-    #    and `resized_image` (for actual tensor conversion).
+    # 3) Run inference
     try:
-        coordinates_str = run_inference_localization(messages, resized_image)
+        coordinates_str = run_inference_localization(messages, resized_image, device, dtype)
     except Exception as e:
-        print(f"Error during model inference: {e}")
-        return f"Error during model inference: {e}", resized_image.copy().convert("RGB")
-
-    # 4. Parse coordinates and draw on the image
-    output_image_with_click = resized_image.copy().convert("RGB") # Ensure it's RGB for drawing
-    parsed_coords = None
-    
-    # Expected format from the model: "Click(x, y)"
+        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))
-            parsed_coords = (x, y)
-            
             draw = ImageDraw.Draw(output_image_with_click)
-            # Make the marker somewhat responsive to image size, but not too small/large
-            radius = max(5, min(resized_width // 100, resized_height // 100, 15)) 
-            
-            # Define the bounding box for the ellipse (circle)
+            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)) # Draw a red circle
+            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 ValueError:
-            print(f"Could not parse integers from coordinates: {coordinates_str}")
-            # Keep original coordinates_str, output_image_with_click will be the resized image without a mark
         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
+
+    return coordinates_str, output_image_with_click, f"device: {device} | dtype: {str(dtype).replace('torch.', '')}"
 
 # --- Load Example Data ---
 example_image = None
-example_instruction = "Select July 14th as the check-out date"
+example_instruction = "Enter the server address readyforquantum.com to check its security"
 try:
-    example_image_url = "https://huggingface.co/Hcompany/Holo1-7B/resolve/main/calendar_example.jpg"
+    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}")
-    # Create a placeholder image if loading fails, so Gradio example still works
+    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: # If PIL itself is an issue (unlikely here but good for robustness)
-        pass 
-
-
-# --- Gradio Interface Definition ---
-title = "Holo1-7B: Action VLM Localization Demo"
-description = """
-This demo showcases **Holo1-7B**, an Action Vision-Language Model developed by HCompany, fine-tuned from Qwen/Qwen2.5-VL-7B-Instruct.
-It's designed to interact with web interfaces like a human user. Here, we demonstrate its UI localization capability.
+    except Exception:
+        pass
 
-**How to use:**
-1. Upload an image (e.g., a screenshot of a UI, like the calendar example).
-2. Provide a textual instruction (e.g., "Select July 14th as the check-out date").
-3. The model will predict the click coordinates in the format `Click(x, y)`.
-4. The predicted click point will be marked with a red circle on the (resized) image.
-
-The model processes a resized version of your input image. Coordinates are relative to this resized image.
-"""
+# --- Gradio UI ---
+title = "Holo1-3B: Holo1 Localization Demo (ZeroGPU-ready)"
 article = f"""
 <p style='text-align: center'>
-Model: <a href='https://huggingface.co/{MODEL_ID}' target='_blank'>{MODEL_ID}</a> by HCompany | 
+Model: <a href='https://huggingface.co/{MODEL_ID}' target='_blank'>{MODEL_ID}</a> by HCompany |
 Paper: <a href='https://cdn.prod.website-files.com/67e2dbd9acff0c50d4c8a80c/683ec8095b353e8b38317f80_h_tech_report_v1.pdf' target='_blank'>HCompany Tech Report</a> |
-Blog: <a href='https://www.hcompany.ai/surfer-h' target='_blank'>Surfer-H Blog Post</a>
+Blog: <a href='https://www.hcompany.ai/surfer-h' target='_blank'>Surfer-H Blog Post</a><br/>
+<small>GPU (if available) is requested only during inference via @spaces.GPU.</small>
 </p>
 """
 
@@ -228,42 +278,54 @@ if not model_loaded:
     with gr.Blocks() as demo:
         gr.Markdown(f"# <center>⚠️ Error: Model Failed to Load ⚠️</center>")
         gr.Markdown(f"<center>{load_error_message}</center>")
-        gr.Markdown("<center>Please check the console output for more details. Reloading the space might help if it's a temporary issue.</center>")
+        gr.Markdown("<center>See logs for the full traceback.</center>")
 else:
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
-        # gr.Markdown(description)
+        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", 
+                    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)
-        
+                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],
+                outputs=[output_coords_component, output_image_component, runtime_info],
                 fn=predict_click_location,
                 cache_examples="lazy",
             )
-        
-        gr.Markdown(article)
 
         submit_button.click(
             fn=predict_click_location,
             inputs=[input_image_component, instruction_component],
-            outputs=[output_coords_component, output_image_component]
+            outputs=[output_coords_component, output_image_component, runtime_info]
         )
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    # Do NOT pass 'concurrency_count' or ZeroGPU-specific launch args.
+    demo.launch(debug=True)