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