Commit
·
9792fe2
1
Parent(s):
deab97e
README.md
CHANGED
@@ -1,8 +1,8 @@
|
|
1 |
---
|
2 |
-
title:
|
3 |
-
emoji:
|
4 |
-
colorFrom:
|
5 |
-
colorTo:
|
6 |
sdk: gradio
|
7 |
sdk_version: 5.42.0
|
8 |
app_file: app.py
|
|
|
1 |
---
|
2 |
+
title: TestHolo
|
3 |
+
emoji: 📊
|
4 |
+
colorFrom: green
|
5 |
+
colorTo: blue
|
6 |
sdk: gradio
|
7 |
sdk_version: 5.42.0
|
8 |
app_file: app.py
|
app.py
ADDED
@@ -0,0 +1,327 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# requirements.txt stays fine, but for CUDA wheels you usually want:
|
2 |
+
# pip install --index-url https://download.pytorch.org/whl/cu121 torch torchvision --upgrade
|
3 |
+
|
4 |
+
import gradio as gr
|
5 |
+
import json, os, re, traceback
|
6 |
+
from typing import Any, List, Dict
|
7 |
+
|
8 |
+
import spaces
|
9 |
+
import torch
|
10 |
+
from PIL import Image, ImageDraw
|
11 |
+
import requests
|
12 |
+
from transformers import AutoModelForImageTextToText, AutoProcessor
|
13 |
+
from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
|
14 |
+
|
15 |
+
# --- Configuration ---
|
16 |
+
MODEL_ID = "Hcompany/Holo1-3B"
|
17 |
+
|
18 |
+
# ---------------- Device / DType helpers ----------------
|
19 |
+
|
20 |
+
def pick_device() -> str:
|
21 |
+
forced = os.getenv("FORCE_DEVICE", "").lower().strip()
|
22 |
+
if forced in {"cpu", "cuda", "mps"}:
|
23 |
+
return forced
|
24 |
+
if torch.cuda.is_available():
|
25 |
+
return "cuda"
|
26 |
+
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
|
27 |
+
return "mps"
|
28 |
+
return "cpu"
|
29 |
+
|
30 |
+
def pick_dtype(device: str) -> torch.dtype:
|
31 |
+
if device == "cuda":
|
32 |
+
major, minor = torch.cuda.get_device_capability() # e.g. (8, 0) for A100
|
33 |
+
# Prefer bfloat16 on Ampere+ (>= 8.x). Otherwise float16.
|
34 |
+
return torch.bfloat16 if major >= 8 else torch.float16
|
35 |
+
if device == "mps":
|
36 |
+
# MPS autocast supports float16 well; bfloat16 is improving but use float16 for safety.
|
37 |
+
return torch.float16
|
38 |
+
return torch.float32 # CPU
|
39 |
+
|
40 |
+
def move_to_device(batch, device: str):
|
41 |
+
if isinstance(batch, dict):
|
42 |
+
return {k: (v.to(device, non_blocking=True) if hasattr(v, "to") else v) for k, v in batch.items()}
|
43 |
+
if hasattr(batch, "to"):
|
44 |
+
return batch.to(device, non_blocking=True)
|
45 |
+
return batch
|
46 |
+
|
47 |
+
# --- Chat/template helpers (unchanged except minor tidy) ---
|
48 |
+
def apply_chat_template_compat(processor, messages: List[Dict[str, Any]]) -> str:
|
49 |
+
tok = getattr(processor, "tokenizer", None)
|
50 |
+
if hasattr(processor, "apply_chat_template"):
|
51 |
+
return processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
52 |
+
if tok is not None and hasattr(tok, "apply_chat_template"):
|
53 |
+
return tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
54 |
+
texts = []
|
55 |
+
for m in messages:
|
56 |
+
for c in m.get("content", []):
|
57 |
+
if isinstance(c, dict) and c.get("type") == "text":
|
58 |
+
texts.append(c.get("text", ""))
|
59 |
+
return "\n".join(texts)
|
60 |
+
|
61 |
+
def batch_decode_compat(processor, token_id_batches, **kw):
|
62 |
+
tok = getattr(processor, "tokenizer", None)
|
63 |
+
if tok is not None and hasattr(tok, "batch_decode"):
|
64 |
+
return tok.batch_decode(token_id_batches, **kw)
|
65 |
+
if hasattr(processor, "batch_decode"):
|
66 |
+
return processor.batch_decode(token_id_batches, **kw)
|
67 |
+
raise AttributeError("No batch_decode available on processor or tokenizer.")
|
68 |
+
|
69 |
+
def get_image_proc_params(processor) -> Dict[str, int]:
|
70 |
+
ip = getattr(processor, "image_processor", None)
|
71 |
+
return {
|
72 |
+
"patch_size": getattr(ip, "patch_size", 14),
|
73 |
+
"merge_size": getattr(ip, "merge_size", 1),
|
74 |
+
"min_pixels": getattr(ip, "min_pixels", 256 * 256),
|
75 |
+
"max_pixels": getattr(ip, "max_pixels", 1280 * 1280),
|
76 |
+
}
|
77 |
+
|
78 |
+
def trim_generated(generated_ids, inputs):
|
79 |
+
in_ids = getattr(inputs, "input_ids", None)
|
80 |
+
if in_ids is None and isinstance(inputs, dict):
|
81 |
+
in_ids = inputs.get("input_ids", None)
|
82 |
+
if in_ids is None:
|
83 |
+
return [out_ids for out_ids in generated_ids]
|
84 |
+
return [out_ids[len(in_seq):] for in_seq, out_ids in zip(in_ids, generated_ids)]
|
85 |
+
|
86 |
+
# --- Load model/processor once with correct device/dtype ---
|
87 |
+
active_device = pick_device()
|
88 |
+
active_dtype = pick_dtype(active_device)
|
89 |
+
|
90 |
+
# Optional perf knobs for CUDA
|
91 |
+
if active_device == "cuda":
|
92 |
+
torch.backends.cuda.matmul.allow_tf32 = True
|
93 |
+
torch.set_float32_matmul_precision("high") # better perf on Ampere+
|
94 |
+
|
95 |
+
print(f"Loading model and processor for {MODEL_ID} on device={active_device}, dtype={active_dtype}...")
|
96 |
+
model = None
|
97 |
+
processor = None
|
98 |
+
model_loaded = False
|
99 |
+
load_error_message = ""
|
100 |
+
|
101 |
+
try:
|
102 |
+
# Note: for single-GPU we explicitly set dtype then .to(device).
|
103 |
+
# If you want HF Accelerate sharding: set device_map="auto" and drop explicit .to().
|
104 |
+
model = AutoModelForImageTextToText.from_pretrained(
|
105 |
+
MODEL_ID,
|
106 |
+
torch_dtype=active_dtype if active_device != "cpu" else torch.float32,
|
107 |
+
trust_remote_code=True,
|
108 |
+
)
|
109 |
+
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
|
110 |
+
|
111 |
+
# Move model to device and eval
|
112 |
+
model.to(active_device)
|
113 |
+
model.eval()
|
114 |
+
model_loaded = True
|
115 |
+
print("Model and processor loaded successfully.")
|
116 |
+
except Exception as e:
|
117 |
+
load_error_message = (
|
118 |
+
f"Error loading model/processor: {e}\n"
|
119 |
+
"This might be due to CUDA/MPS availability, model ID, or wheel incompatibility.\n"
|
120 |
+
"Check the full traceback in the logs."
|
121 |
+
)
|
122 |
+
print(load_error_message)
|
123 |
+
traceback.print_exc()
|
124 |
+
|
125 |
+
# --- Prompt builder ---
|
126 |
+
def get_localization_prompt(pil_image: Image.Image, instruction: str) -> List[dict]:
|
127 |
+
guidelines: str = (
|
128 |
+
"Localize an element on the GUI image according to my instructions and "
|
129 |
+
"output a click position as Click(x, y) with x num pixels from the left edge "
|
130 |
+
"and y num pixels from the top edge."
|
131 |
+
)
|
132 |
+
return [
|
133 |
+
{
|
134 |
+
"role": "user",
|
135 |
+
"content": [
|
136 |
+
{"type": "image", "image": pil_image},
|
137 |
+
{"type": "text", "text": f"{guidelines}\n{instruction}"}
|
138 |
+
],
|
139 |
+
}
|
140 |
+
]
|
141 |
+
|
142 |
+
# --- Inference (device-agnostic; uses AMP on GPU) ---
|
143 |
+
@torch.inference_mode()
|
144 |
+
def run_inference_localization(
|
145 |
+
messages_for_template: List[dict[str, Any]],
|
146 |
+
pil_image_for_processing: Image.Image
|
147 |
+
) -> str:
|
148 |
+
try:
|
149 |
+
# 1) Build prompt text
|
150 |
+
text_prompt = apply_chat_template_compat(processor, messages_for_template)
|
151 |
+
|
152 |
+
# 2) Prepare inputs (text + image)
|
153 |
+
inputs = processor(
|
154 |
+
text=[text_prompt],
|
155 |
+
images=[pil_image_for_processing],
|
156 |
+
padding=True,
|
157 |
+
return_tensors="pt",
|
158 |
+
)
|
159 |
+
inputs = move_to_device(inputs, active_device)
|
160 |
+
|
161 |
+
# 3) Generate (deterministic). Use autocast on GPU/MPS.
|
162 |
+
use_amp = active_device in {"cuda", "mps"}
|
163 |
+
amp_dtype = active_dtype if active_device == "cuda" else torch.float16
|
164 |
+
|
165 |
+
if use_amp:
|
166 |
+
with torch.cuda.amp.autocast(enabled=(active_device == "cuda"), dtype=amp_dtype):
|
167 |
+
generated_ids = model.generate(
|
168 |
+
**inputs,
|
169 |
+
max_new_tokens=128,
|
170 |
+
do_sample=False,
|
171 |
+
)
|
172 |
+
else:
|
173 |
+
generated_ids = model.generate(
|
174 |
+
**inputs,
|
175 |
+
max_new_tokens=128,
|
176 |
+
do_sample=False,
|
177 |
+
)
|
178 |
+
|
179 |
+
# 4) Trim prompt tokens if possible
|
180 |
+
generated_ids_trimmed = trim_generated(generated_ids, inputs)
|
181 |
+
|
182 |
+
# 5) Decode
|
183 |
+
decoded_output = batch_decode_compat(
|
184 |
+
processor,
|
185 |
+
generated_ids_trimmed,
|
186 |
+
skip_special_tokens=True,
|
187 |
+
clean_up_tokenization_spaces=False
|
188 |
+
)
|
189 |
+
|
190 |
+
return decoded_output[0] if decoded_output else ""
|
191 |
+
except Exception as e:
|
192 |
+
print(f"Error during model inference: {e}")
|
193 |
+
traceback.print_exc()
|
194 |
+
raise
|
195 |
+
|
196 |
+
# --- Gradio processing function ---
|
197 |
+
def predict_click_location(input_pil_image: Image.Image, instruction: str):
|
198 |
+
if not model_loaded or not processor or not model:
|
199 |
+
return f"Model not loaded. Error: {load_error_message}", None
|
200 |
+
if not input_pil_image:
|
201 |
+
return "No image provided. Please upload an image.", None
|
202 |
+
if not instruction or instruction.strip() == "":
|
203 |
+
return "No instruction provided. Please type an instruction.", input_pil_image.copy().convert("RGB")
|
204 |
+
|
205 |
+
# 1) Resize according to image processor params (safe defaults if missing)
|
206 |
+
try:
|
207 |
+
ip = get_image_proc_params(processor)
|
208 |
+
resized_height, resized_width = smart_resize(
|
209 |
+
input_pil_image.height,
|
210 |
+
input_pil_image.width,
|
211 |
+
factor=ip["patch_size"] * ip["merge_size"],
|
212 |
+
min_pixels=ip["min_pixels"],
|
213 |
+
max_pixels=ip["max_pixels"],
|
214 |
+
)
|
215 |
+
resized_image = input_pil_image.resize(
|
216 |
+
size=(resized_width, resized_height),
|
217 |
+
resample=Image.Resampling.LANCZOS
|
218 |
+
)
|
219 |
+
except Exception as e:
|
220 |
+
print(f"Error resizing image: {e}")
|
221 |
+
traceback.print_exc()
|
222 |
+
return f"Error resizing image: {e}", input_pil_image.copy().convert("RGB")
|
223 |
+
|
224 |
+
# 2) Build messages with image + instruction
|
225 |
+
messages = get_localization_prompt(resized_image, instruction)
|
226 |
+
|
227 |
+
# 3) Run inference
|
228 |
+
try:
|
229 |
+
coordinates_str = run_inference_localization(messages, resized_image)
|
230 |
+
except Exception as e:
|
231 |
+
return f"Error during model inference: {e}", resized_image.copy().convert("RGB")
|
232 |
+
|
233 |
+
# 4) Parse coordinates and draw marker
|
234 |
+
output_image_with_click = resized_image.copy().convert("RGB")
|
235 |
+
match = re.search(r"Click\((\d+),\s*(\d+)\)", coordinates_str)
|
236 |
+
if match:
|
237 |
+
try:
|
238 |
+
x = int(match.group(1)); y = int(match.group(2))
|
239 |
+
draw = ImageDraw.Draw(output_image_with_click)
|
240 |
+
radius = max(5, min(resized_width // 100, resized_height // 100, 15))
|
241 |
+
bbox = (x - radius, y - radius, x + radius, y + radius)
|
242 |
+
draw.ellipse(bbox, outline="red", width=max(2, radius // 4))
|
243 |
+
print(f"Predicted and drawn click at: ({x}, {y}) on resized image ({resized_width}x{resized_height})")
|
244 |
+
except Exception as e:
|
245 |
+
print(f"Error drawing on image: {e}")
|
246 |
+
traceback.print_exc()
|
247 |
+
else:
|
248 |
+
print(f"Could not parse 'Click(x, y)' from model output: {coordinates_str}")
|
249 |
+
|
250 |
+
return coordinates_str, output_image_with_click
|
251 |
+
|
252 |
+
# --- Load Example Data ---
|
253 |
+
example_image = None
|
254 |
+
example_instruction = "Enter the server address readyforquantum.com to check its security"
|
255 |
+
try:
|
256 |
+
example_image_url = "https://readyforquantum.com/img/screentest.jpg"
|
257 |
+
example_image = Image.open(requests.get(example_image_url, stream=True).raw)
|
258 |
+
except Exception as e:
|
259 |
+
print(f"Could not load example image from URL: {e}")
|
260 |
+
traceback.print_exc()
|
261 |
+
try:
|
262 |
+
example_image = Image.new("RGB", (200, 150), color="lightgray")
|
263 |
+
draw = ImageDraw.Draw(example_image)
|
264 |
+
draw.text((10, 10), "Example image\nfailed to load", fill="black")
|
265 |
+
except Exception:
|
266 |
+
pass
|
267 |
+
|
268 |
+
# --- Gradio UI ---
|
269 |
+
title = "Holo1-3B: Holo1 Localization Demo"
|
270 |
+
article = f"""
|
271 |
+
<p style='text-align: center'>
|
272 |
+
Device: <b>{active_device}</b> | DType: <b>{str(active_dtype).replace('torch.', '')}</b> |
|
273 |
+
Model: <a href='https://huggingface.co/{MODEL_ID}' target='_blank'>{MODEL_ID}</a> by HCompany |
|
274 |
+
Paper: <a href='https://cdn.prod.website-files.com/67e2dbd9acff0c50d4c8a80c/683ec8095b353e8b38317f80_h_tech_report_v1.pdf' target='_blank'>HCompany Tech Report</a> |
|
275 |
+
Blog: <a href='https://www.hcompany.ai/surfer-h' target='_blank'>Surfer-H Blog Post</a>
|
276 |
+
</p>
|
277 |
+
"""
|
278 |
+
|
279 |
+
if not model_loaded:
|
280 |
+
with gr.Blocks() as demo:
|
281 |
+
gr.Markdown(f"# <center>⚠️ Error: Model Failed to Load ⚠️</center>")
|
282 |
+
gr.Markdown(f"<center>{load_error_message}</center>")
|
283 |
+
gr.Markdown("<center>See logs for the full traceback.</center>")
|
284 |
+
else:
|
285 |
+
with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
286 |
+
gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
|
287 |
+
gr.Markdown(article)
|
288 |
+
|
289 |
+
with gr.Row():
|
290 |
+
with gr.Column(scale=1):
|
291 |
+
input_image_component = gr.Image(type="pil", label="Input UI Image", height=400)
|
292 |
+
instruction_component = gr.Textbox(
|
293 |
+
label="Instruction",
|
294 |
+
placeholder="e.g., Click the 'Login' button",
|
295 |
+
info="Type the action you want the model to localize on the image."
|
296 |
+
)
|
297 |
+
submit_button = gr.Button("Localize Click", variant="primary")
|
298 |
+
|
299 |
+
with gr.Column(scale=1):
|
300 |
+
output_coords_component = gr.Textbox(
|
301 |
+
label="Predicted Coordinates (Format: Click(x, y))",
|
302 |
+
interactive=False
|
303 |
+
)
|
304 |
+
output_image_component = gr.Image(
|
305 |
+
type="pil",
|
306 |
+
label="Image with Predicted Click Point",
|
307 |
+
height=400,
|
308 |
+
interactive=False
|
309 |
+
)
|
310 |
+
|
311 |
+
if example_image:
|
312 |
+
gr.Examples(
|
313 |
+
examples=[[example_image, example_instruction]],
|
314 |
+
inputs=[input_image_component, instruction_component],
|
315 |
+
outputs=[output_coords_component, output_image_component],
|
316 |
+
fn=predict_click_location,
|
317 |
+
cache_examples="lazy",
|
318 |
+
)
|
319 |
+
|
320 |
+
submit_button.click(
|
321 |
+
fn=predict_click_location,
|
322 |
+
inputs=[input_image_component, instruction_component],
|
323 |
+
outputs=[output_coords_component, output_image_component]
|
324 |
+
)
|
325 |
+
|
326 |
+
if __name__ == "__main__":
|
327 |
+
demo.launch(debug=True)
|
commit
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
git add .
|
2 |
+
git commit -m "$*"
|
3 |
+
git push
|
requirements.txt
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
transformers
|
2 |
+
accelerate
|
3 |
+
torch
|
4 |
+
torchvision
|
5 |
+
gradio
|
6 |
+
spaces
|
7 |
+
Pillow
|
8 |
+
requests
|
9 |
+
|