Update app.py

app.py

@@ -29,9 +29,22 @@ from transformers.image_utils import load_image
 # --- Constants and Model Setup ---
 MAX_INPUT_TOKEN_LENGTH = 4096
 # Note: The following line correctly falls back to CPU if CUDA is not available.
-# The "RuntimeError: CUDA driver initialization failed" is an environment issue,
-# meaning the code is being run where a GPU is expected but not found/configured.
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Let the environment (e.g., Hugging Face Spaces) determine the device.
+# This avoids conflicts with the CUDA environment setup by the platform.
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
+print("torch.__version__ =", torch.__version__)
+print("torch.version.cuda =", torch.version.cuda)
+print("cuda available:", torch.cuda.is_available())
+print("cuda device count:", torch.cuda.device_count())
+if torch.cuda.is_available():
+    print("current device:", torch.cuda.current_device())
+    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
+
+print("Using device:", device)
+
+# --- Model Loading ---
 
 # --- Prompts for Different Tasks ---
 layout_prompt = """Please output the layout information from the image, including each layout element's bbox, its category, and the corresponding text content within the bbox.