Update app.py

app.py

@@ -1,8 +1,5 @@
-# requirements.txt stays fine, but for CUDA wheels you usually want:
-# pip install --index-url https://download.pytorch.org/whl/cu121 torch torchvision --upgrade
-
 import gradio as gr
-import json, os, re, traceback
+import json, os, re, traceback, contextlib
 from typing import Any, List, Dict
 
 import spaces
@@ -18,6 +15,10 @@ 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
@@ -29,11 +30,9 @@ def pick_device() -> str:
 
 def pick_dtype(device: str) -> torch.dtype:
     if device == "cuda":
-        major, minor = torch.cuda.get_device_capability()  # e.g. (8, 0) for A100
-        # Prefer bfloat16 on Ampere+ (>= 8.x). Otherwise float16.
-        return torch.bfloat16 if major >= 8 else torch.float16
+        major, _ = torch.cuda.get_device_capability()
+        return torch.bfloat16 if major >= 8 else torch.float16  # Ampere+ -> bf16
     if device == "mps":
-        # MPS autocast supports float16 well; bfloat16 is improving but use float16 for safety.
         return torch.float16
     return torch.float32  # CPU
 
@@ -44,7 +43,7 @@ def move_to_device(batch, device: str):
         return batch.to(device, non_blocking=True)
     return batch
 
-# --- Chat/template helpers (unchanged except minor tidy) ---
+# --- 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"):
@@ -83,40 +82,27 @@ def trim_generated(generated_ids, inputs):
         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 once with correct device/dtype ---
-active_device = pick_device()
-active_dtype = pick_dtype(active_device)
-
-# Optional perf knobs for CUDA
-if active_device == "cuda":
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.set_float32_matmul_precision("high")  # better perf on Ampere+
-
-print(f"Loading model and processor for {MODEL_ID} on device={active_device}, dtype={active_dtype}...")
+# --- 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:
-    # Note: for single-GPU we explicitly set dtype then .to(device).
-    # If you want HF Accelerate sharding: set device_map="auto" and drop explicit .to().
     model = AutoModelForImageTextToText.from_pretrained(
         MODEL_ID,
-        torch_dtype=active_dtype if active_device != "cpu" else torch.float32,
+        torch_dtype=torch.float32,  # CPU-safe dtype at import
         trust_remote_code=True,
     )
     processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-    # Move model to device and eval
-    model.to(active_device)
     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 CUDA/MPS availability, model ID, or wheel incompatibility.\n"
+        "This might be due to network/model ID/library versions.\n"
         "Check the full traceback in the logs."
     )
     print(load_error_message)
@@ -139,68 +125,79 @@ def get_localization_prompt(pil_image: Image.Image, instruction: str) -> List[di
         }
     ]
 
-# --- Inference (device-agnostic; uses AMP on GPU) ---
+# --- 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
+    pil_image_for_processing: Image.Image,
+    device: str,
+    dtype: torch.dtype,
 ) -> str:
-    try:
-        # 1) Build prompt text
-        text_prompt = apply_chat_template_compat(processor, messages_for_template)
-
-        # 2) Prepare inputs (text + image)
-        inputs = processor(
-            text=[text_prompt],
-            images=[pil_image_for_processing],
-            padding=True,
-            return_tensors="pt",
-        )
-        inputs = move_to_device(inputs, active_device)
-
-        # 3) Generate (deterministic). Use autocast on GPU/MPS.
-        use_amp = active_device in {"cuda", "mps"}
-        amp_dtype = active_dtype if active_device == "cuda" else torch.float16
-
-        if use_amp:
-            with torch.cuda.amp.autocast(enabled=(active_device == "cuda"), dtype=amp_dtype):
-                generated_ids = model.generate(
-                    **inputs,
-                    max_new_tokens=128,
-                    do_sample=False,
-                )
-        else:
-            generated_ids = model.generate(
-                **inputs,
-                max_new_tokens=128,
-                do_sample=False,
-            )
+    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)
 
-        # 4) Trim prompt tokens if possible
-        generated_ids_trimmed = trim_generated(generated_ids, inputs)
+    # 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()
 
-        # 5) Decode
-        decoded_output = batch_decode_compat(
-            processor,
-            generated_ids_trimmed,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=False
+    with amp_ctx:
+        generated_ids = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            do_sample=False,
         )
 
-        return decoded_output[0] if decoded_output else ""
-    except Exception as e:
-        print(f"Error during model inference: {e}")
-        traceback.print_exc()
-        raise
+    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"
+
+    # 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:
@@ -217,25 +214,26 @@ def predict_click_location(input_pil_image: Image.Image, instruction: str):
             resample=Image.Resampling.LANCZOS
         )
     except Exception as e:
-        print(f"Error resizing image: {e}")
         traceback.print_exc()
-        return f"Error resizing image: {e}", input_pil_image.copy().convert("RGB")
+        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)
+        coordinates_str = run_inference_localization(messages, resized_image, device, dtype)
     except Exception as e:
-        return f"Error during model inference: {e}", resized_image.copy().convert("RGB")
+        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))
+            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)
@@ -247,7 +245,7 @@ def predict_click_location(input_pil_image: Image.Image, instruction: str):
     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
@@ -266,13 +264,13 @@ except Exception as e:
         pass
 
 # --- Gradio UI ---
-title = "Holo1-3B: Holo1 Localization Demo"
+title = "Holo1-3B: Holo1 Localization Demo (ZeroGPU-ready)"
 article = f"""
 <p style='text-align: center'>
-Device: <b>{active_device}</b> &nbsp;|&nbsp; DType: <b>{str(active_dtype).replace('torch.', '')}</b> &nbsp;|&nbsp;
 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>
 """
 
@@ -307,12 +305,17 @@ else:
                     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",
             )
@@ -320,8 +323,9 @@ else:
         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__":
+    # Do NOT pass 'concurrency_count' or ZeroGPU-specific launch args.
     demo.launch(debug=True)