get more info HF space

app.py

@@ -12,6 +12,8 @@ from zea import init_device
 
 from main import Config, init, run
 from utils import load_image
+import torch
+import subprocess
 
 CONFIG_PATH = "configs/semantic_dps.yaml"
 SLIDER_CONFIG_PATH = "configs/slider_params.yaml"
@@ -285,9 +287,30 @@ with gr.Blocks() as demo:
                     DEVICE = init_device()
                 except:
                     print("Could not initialize device using `zea.init_device()`")
-                    print(f"JAX version: {jax.__version__}")
-                    print(f"JAX devices: {jax.devices()}")
                     print(f"KERAS version: {keras.__version__}")
+                    try:
+                        print(f"JAX version: {jax.__version__}")
+                        print(f"JAX devices: {jax.devices()}")
+                    except Exception as e:
+                        print(f"Could not get JAX info: {e}")
+
+                    try:
+                        print(f"PyTorch version: {torch.__version__}")
+                        print(f"PyTorch CUDA available: {torch.cuda.is_available()}")
+                        print(f"PyTorch CUDA device count: {torch.cuda.device_count()}")
+                        print(f"PyTorch devices: {[torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]}")
+                        print(f"PyTorch CUDA version: {torch.version.cuda}")
+                        print(f"PyTorch cuDNN version: {torch.backends.cudnn.version()}")
+                    except Exception as e:
+                        print(f"Could not get PyTorch info: {e}")
+
+                    try:
+                        cuda_version = subprocess.getoutput("nvcc --version")
+                        print(f"nvcc version:\n{cuda_version}")
+                        nvidia_smi = subprocess.getoutput("nvidia-smi")
+                        print(f"nvidia-smi output:\n{nvidia_smi}")
+                    except Exception as e:
+                        print(f"Could not get CUDA/nvidia-smi info: {e}")
 
             config, diffusion_model = initialize_model()
             ready_msg = gr.update(