johnbridges commited on
Commit
9792fe2
·
1 Parent(s): deab97e
Files changed (4) hide show
  1. README.md +4 -4
  2. app.py +327 -0
  3. commit +3 -0
  4. requirements.txt +9 -0
README.md CHANGED
@@ -1,8 +1,8 @@
1
  ---
2
- title: Holo1
3
- emoji: 🐨
4
- colorFrom: gray
5
- colorTo: gray
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> &nbsp;|&nbsp; DType: <b>{str(active_dtype).replace('torch.', '')}</b> &nbsp;|&nbsp;
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
+