Fix meta-device inference errors

app.py

@@ -37,6 +37,20 @@ processor = None
 inference_device = None
 
 
+def resolve_inference_device(model_obj) -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+
+    # With device_map='auto', some parameters can live on 'meta' while offloaded.
+    # For CPU Spaces, inputs must stay on CPU.
+    device_map = getattr(model_obj, "hf_device_map", None)
+    if isinstance(device_map, dict):
+        for mapped in device_map.values():
+            if isinstance(mapped, str) and mapped.startswith("cuda"):
+                return torch.device(mapped)
+    return torch.device("cpu")
+
+
 def is_tmos_adapter_config(cfg: dict) -> bool:
     modules_to_save = cfg.get("modules_to_save") or []
     target_modules = set(cfg.get("target_modules") or [])
@@ -149,7 +163,7 @@ def load_model_and_processor():
     processor.patch_size = 14
     processor.vision_feature_select_strategy = "default"
 
-    inference_device = next(model.parameters()).device
+    inference_device = resolve_inference_device(model)
     if adapter_source == ADAPTER_REPO_ID:
         print(f"TMOS-DD ready on {inference_device} using remote subfolder '{selected_subfolder or '.'}'.")
     else:
@@ -188,46 +202,51 @@ def confidence_card(prob_fake: float, label: str) -> str:
 
 
 def infer_image(image: Image.Image):
-    if image is None:
-        raise gr.Error("Please upload an image before running detection.")
-
-    model_obj, processor_obj, device = load_model_and_processor()
+    try:
+        if image is None:
+            raise gr.Error("Please upload an image before running detection.")
 
-    prepared_image = preprocess_image(image)
-    inputs = processor_obj(text=PROMPT, images=prepared_image, return_tensors="pt", padding=True)
-    inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
+        model_obj, processor_obj, device = load_model_and_processor()
 
-    autocast_context = (
-        torch.autocast(device_type="cuda", dtype=select_torch_dtype())
-        if device.type == "cuda"
-        else nullcontext()
-    )
+        prepared_image = preprocess_image(image)
+        inputs = processor_obj(text=PROMPT, images=prepared_image, return_tensors="pt", padding=True)
+        inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
 
-    start_time = time.perf_counter()
-    with torch.inference_mode(), autocast_context:
-        outputs = model_obj(
-            input_ids=inputs["input_ids"],
-            pixel_values=inputs["pixel_values"],
-            attention_mask=inputs["attention_mask"],
+        autocast_context = (
+            torch.autocast(device_type="cuda", dtype=select_torch_dtype())
+            if device.type == "cuda"
+            else nullcontext()
         )
 
-    if device.type == "cuda":
-        torch.cuda.synchronize()
+        start_time = time.perf_counter()
+        with torch.inference_mode(), autocast_context:
+            outputs = model_obj(
+                input_ids=inputs["input_ids"],
+                pixel_values=inputs["pixel_values"],
+                attention_mask=inputs["attention_mask"],
+            )
 
-    elapsed_ms = (time.perf_counter() - start_time) * 1000.0
+        if device.type == "cuda":
+            torch.cuda.synchronize()
 
-    logit = float(outputs["logit"].squeeze().detach().float().cpu().item())
-    if not math.isfinite(logit):
-        raise gr.Error("Model produced a non-finite logit (NaN/Inf). Please retry.")
+        elapsed_ms = (time.perf_counter() - start_time) * 1000.0
 
-    prob_fake = float(torch.sigmoid(torch.tensor(logit)).item())
-    if not math.isfinite(prob_fake):
-        raise gr.Error("Model produced a non-finite probability (NaN/Inf). Please retry.")
+        logit = float(outputs["logit"].squeeze().detach().float().cpu().item())
+        if not math.isfinite(logit):
+            raise gr.Error("Model produced a non-finite logit (NaN/Inf). Please retry.")
 
-    label = "Fake" if prob_fake >= THRESHOLD else "Real"
-    confidence = prob_fake if label == "Fake" else 1.0 - prob_fake
+        prob_fake = float(torch.sigmoid(torch.tensor(logit)).item())
+        if not math.isfinite(prob_fake):
+            raise gr.Error("Model produced a non-finite probability (NaN/Inf). Please retry.")
+
+        label = "Fake" if prob_fake >= THRESHOLD else "Real"
+        confidence = prob_fake if label == "Fake" else 1.0 - prob_fake
 
-    return prepared_image, label, round(prob_fake, 6), round(confidence * 100.0, 2), round(elapsed_ms, 2), confidence_card(prob_fake, label)
+        return prepared_image, label, round(prob_fake, 6), round(confidence * 100.0, 2), round(elapsed_ms, 2), confidence_card(prob_fake, label)
+    except gr.Error:
+        raise
+    except Exception as exc:
+        raise gr.Error(f"Inference failed: {type(exc).__name__}: {exc}") from exc
 
 
 load_model_and_processor()