fixed docker issues

.dockerignore

@@ -0,0 +1,17 @@
+.git
+.gitattributes
+.gradio
+.idea
+.vscode
+**/__pycache__
+**/*.py[cod]
+*.egg-info
+**/*.egg-info
+**/.ipynb_checkpoints
+*.ipynb
+results
+.pytest_cache
+.ruff_cache
+.mypy_cache
+.venv
+venv

Dockerfile

@@ -1,41 +1,60 @@
-FROM python:3.10
+# PyTorch 2.9.1 + CUDA 13.0 includes Blackwell (sm_120) support required by
+# GeForce RTX 50-series GPUs. Pin the digest so rebuilds use the same runtime.
+FROM pytorch/pytorch:2.9.1-cuda13.0-cudnn9-runtime@sha256:60f22fb80755fd0b470fb47928dbd55816aa9f847edd95cf43c93253507a9ddf
 
-COPY ./requirements-gradio.txt /code/requirements-gradio.txt
+ARG APP_UID=1000
+ARG APP_GID=1000
 
-# Install system dependencies and create user
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    && useradd -m -u 1000 user \
-    && rm -rf /var/lib/apt/lists/*
-
-# Install OpenGL and other dependencies required for OpenCV
-RUN apt-get update && apt-get install -y \
-    libgl1-mesa-glx \
-    libglib2.0-0 \
-    && rm -rf /var/lib/apt/lists/*
-
-# Switch to "user" before installing dependencies
-USER user
-ENV HOME=/home/user \
-    PATH=/home/user/.local/bin:$PATH \
-    PYTHONPATH=$HOME/app \
+ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    HOME=/home/user \
+    PATH=/home/user/.local/bin:/opt/conda/bin:$PATH \
+    PYTHONPATH=/home/user/app \
+    HF_HOME=/home/user/.cache/huggingface \
+    TORCH_HOME=/home/user/.cache/torch \
+    MPLCONFIGDIR=/home/user/.cache/matplotlib \
     GRADIO_ALLOW_FLAGGING=never \
-    GRADIO_NUM_PORTS=1 \
     GRADIO_SERVER_NAME=0.0.0.0 \
-    GRADIO_THEME=huggingface \
-    SYSTEM=spaces
+    GRADIO_SERVER_PORT=7860 \
+    GRADIO_SHARE=false \
+    REQUIRE_CUDA=1 \
+    PYTORCH_ALLOC_CONF=expandable_segments:True
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends \
+        ca-certificates \
+        fonts-dejavu-core \
+        libglib2.0-0 \
+        libgomp1 \
+    && groupadd --gid "${APP_GID}" user \
+    && useradd --create-home --uid "${APP_UID}" --gid "${APP_GID}" --shell /bin/bash user \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /home/user/app
 
-WORKDIR $HOME/app
+# Install pinned dependencies before copying the source to preserve the build
+# cache when only application code changes.
+COPY requirements.txt requirements-gradio.txt ./
+RUN python -m pip install --no-cache-dir \
+        --requirement requirements.txt \
+        --requirement requirements-gradio.txt
 
-# Copy project files as "user" before installing dependencies
-COPY --chown=user . $HOME/app
-COPY --chown=user ./images /home/user/app/images
+COPY --chown=user:user . .
 
-# Install dependencies as "user"
-RUN pip install --no-cache-dir --user -e .
-RUN pip install --no-cache-dir --user opencv-python
-# RUN pip install --no-cache-dir --user -r /code/requirements-gradio.txt
-RUN ls -l /home/user/app/images/
+RUN mkdir -p \
+        /home/user/.cache/huggingface \
+        /home/user/.cache/torch \
+        /home/user/.cache/matplotlib \
+        /home/user/app/results \
+    && python -c "import cv2, gradio, spaces, torch, torchvision, transformers; assert torch.__version__.startswith('2.9.1'); assert torch.version.cuda == '13.0'; assert transformers.__version__ == '4.48.2'" \
+    && chown -R user:user /home/user/.cache /home/user/app/results
 
-CMD ["python", "app.py"]
+EXPOSE 7860
 
+# Fail fast when the container was started without GPU access, then replace the
+# checker process with the application so signals reach Gradio correctly. The
+# checker starts as root only long enough to prepare mounted volumes, then drops
+# to the unprivileged "user" account before touching models or running Gradio.
+CMD ["python", "docker/verify_runtime.py", "--", "python", "app.py"]

README.md

@@ -10,4 +10,31 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at <https://huggingface.co/docs/hub/spaces-config-reference>
+Check out the configuration reference at <https://huggingface.co/docs/hub/spaces-config-reference>
+
+## Docker GPU runtime
+
+The container is pinned to PyTorch 2.9.1 with CUDA 13.0 for NVIDIA
+Blackwell/RTX 50-series compatibility.
+
+Requirements:
+
+- NVIDIA driver compatible with CUDA 13.0
+- Docker Desktop with the WSL2 backend on Windows
+- NVIDIA GPU access enabled in Docker
+
+Build and run:
+
+```shell
+docker compose build
+docker compose up
+```
+
+Open <http://localhost:7860>. The entry point executes a CUDA kernel before
+starting Gradio and exits with an actionable error if the GPU is unavailable.
+
+To validate Docker GPU passthrough independently:
+
+```shell
+docker run --rm --gpus all nvidia/cuda:13.0.0-base-ubuntu22.04 nvidia-smi
+```

app.py

@@ -330,5 +330,9 @@ with gr.Blocks() as demo:
         outputs=[understanding_output, activation_map_output, understanding_target_token_decoded_output]
     )
     
-demo.launch(share=True)
-# demo.queue(concurrency_count=1, max_size=10).launch(server_name="0.0.0.0", server_port=37906, root_path="/path")
+demo.launch(
+    server_name=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
+    server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
+    share=os.getenv("GRADIO_SHARE", "false").lower() in {"1", "true", "yes", "on"},
+)
+# demo.queue(concurrency_count=1, max_size=10).launch(server_name="0.0.0.0", server_port=37906, root_path="/path")

compose.yaml

@@ -0,0 +1,28 @@
+services:
+  app:
+    image: probing-vis-literacy:gpu
+    build:
+      context: .
+      dockerfile: Dockerfile
+    ports:
+      - "7860:7860"
+    environment:
+      REQUIRE_CUDA: "1"
+      GRADIO_SHARE: "false"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    shm_size: "8gb"
+    volumes:
+      - huggingface-cache:/home/user/.cache/huggingface
+      - model-cache:/home/user/.cache/torch
+      - app-results:/home/user/app/results
+
+volumes:
+  huggingface-cache:
+  model-cache:
+  app-results:

demo/model_utils.py

@@ -6,17 +6,30 @@ from transformers import AutoConfig, AutoModelForCausalLM, LlavaForConditionalGe
 from transformers import CLIPProcessor, CLIPModel
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 
+def get_inference_dtype(precision=16):
+    cuda_available = torch.cuda.is_available()
+    if precision == 16 and cuda_available:
+        return (
+            torch.bfloat16
+            if torch.cuda.is_bf16_supported()
+            else torch.float16
+        )
+
+    # Float16 model execution is not consistently supported on CPU.
+    return torch.float32
+
+
 @spaces.GPU(duration=120)
 def set_dtype_device(model, precision=16, device_map=None):
-    dtype = (torch.bfloat16 if torch.cuda.is_available() else torch.float16) if precision==16 else (torch.bfloat32 if torch.cuda.is_available() else torch.float32)
-    cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    if torch.cuda.is_available():
-        model = model.to(dtype)
-        if not device_map:
-            model.cuda()
-    else:
-        torch.set_default_device("cpu")
-        model = model.to(dtype)
+    dtype = get_inference_dtype(precision)
+    cuda_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if device_map:
+        # Accelerate owns device placement for dispatched models. Calling
+        # model.cuda() or model.to(device) afterwards can invalidate its hooks.
+        return model, dtype, cuda_device
+
+    model = model.to(device=cuda_device, dtype=dtype)
     return model, dtype, cuda_device
 
 
@@ -64,6 +77,8 @@ class Janus_Utils(Model_Utils):
         language_config._attn_implementation = 'eager'
         self.vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
                                                     language_config=language_config,
+                                                    torch_dtype=get_inference_dtype(),
+                                                    low_cpu_mem_usage=True,
                                                     trust_remote_code=True,
                                                     ignore_mismatched_sizes=True,
                                                     )
@@ -127,10 +142,13 @@ class LLaVA_Utils(Model_Utils):
             self.vl_gpt = LlavaForConditionalGeneration.from_pretrained(model_path,
                                                         low_cpu_mem_usage=True,
                                                         attn_implementation = 'eager',
+                                                        torch_dtype=get_inference_dtype(),
                                                         device_map="auto",
                                                         output_attentions=True
                                                         )
-            self.vl_gpt, self.dtype, self.cuda_device = set_dtype_device(self.vl_gpt)
+            self.vl_gpt, self.dtype, self.cuda_device = set_dtype_device(
+                self.vl_gpt, device_map="auto"
+            )
             self.processor = AutoProcessor.from_pretrained(model_path)
             self.tokenizer = self.processor.tokenizer
 
@@ -140,7 +158,7 @@ class LLaVA_Utils(Model_Utils):
             self.processor = AutoProcessor.from_pretrained(model_path)
 
             self.vl_gpt = LlavaOnevisionForConditionalGeneration.from_pretrained(model_path, 
-                                                                        torch_dtype=torch.float16,
+                                                                        torch_dtype=get_inference_dtype(),
                                                                         device_map="auto", 
                                                                         low_cpu_mem_usage=True,
                                                                         attn_implementation = 'eager',
@@ -220,7 +238,8 @@ class ChartGemma_Utils(Model_Utils):
         
         self.vl_gpt = PaliGemmaForConditionalGeneration.from_pretrained(
             model_path, 
-            torch_dtype=torch.float16,  
+            torch_dtype=get_inference_dtype(),
+            low_cpu_mem_usage=True,
             attn_implementation="eager", 
             output_attentions=True
         )
@@ -294,4 +313,3 @@ def add_title_to_image(image, title, font_size=50):
     return combined
     
 
-    

docker/verify_runtime.py

@@ -0,0 +1,104 @@
+"""Fail-fast CUDA validation for the Docker entry point."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+import pwd
+import sys
+
+import torch
+
+
+WRITABLE_DIRECTORIES = (
+    Path("/home/user/.cache/huggingface"),
+    Path("/home/user/.cache/torch"),
+    Path("/home/user/.cache/matplotlib"),
+    Path("/home/user/app/results"),
+)
+
+
+def env_flag(name: str, default: bool) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def prepare_directories_and_drop_privileges() -> None:
+    """Make mounted volumes writable, then permanently become the app user."""
+    if os.geteuid() != 0:
+        return
+
+    account = pwd.getpwnam("user")
+    for directory in WRITABLE_DIRECTORIES:
+        directory.mkdir(parents=True, exist_ok=True)
+        for root, directories, files in os.walk(directory):
+            os.chown(root, account.pw_uid, account.pw_gid)
+            for name in directories:
+                os.chown(
+                    os.path.join(root, name),
+                    account.pw_uid,
+                    account.pw_gid,
+                    follow_symlinks=False,
+                )
+            for name in files:
+                os.chown(
+                    os.path.join(root, name),
+                    account.pw_uid,
+                    account.pw_gid,
+                    follow_symlinks=False,
+                )
+
+    os.setgroups([])
+    os.setgid(account.pw_gid)
+    os.setuid(account.pw_uid)
+
+    if os.geteuid() == 0:
+        raise RuntimeError("Failed to drop root privileges.")
+
+
+def verify_cuda() -> None:
+    require_cuda = env_flag("REQUIRE_CUDA", True)
+
+    if not torch.cuda.is_available():
+        message = (
+            "CUDA is not visible inside the container. Start it with "
+            "`docker compose up` or `docker run --gpus all ...`."
+        )
+        if require_cuda:
+            raise RuntimeError(message)
+        print(f"WARNING: {message} Continuing because REQUIRE_CUDA=0.", flush=True)
+        return
+
+    # Executing a real kernel detects runtime/driver/architecture mismatches that
+    # torch.cuda.is_available() alone can miss.
+    device = torch.device("cuda:0")
+    result = (torch.ones(1, device=device) + 1).item()
+    torch.cuda.synchronize(device)
+    if result != 2:
+        raise RuntimeError("CUDA smoke test returned an unexpected result.")
+
+    properties = torch.cuda.get_device_properties(device)
+    capability = torch.cuda.get_device_capability(device)
+    print(
+        "CUDA ready: "
+        f"torch={torch.__version__}, runtime={torch.version.cuda}, "
+        f"device={properties.name}, capability=sm_{capability[0]}{capability[1]}",
+        flush=True,
+    )
+
+
+def main() -> None:
+    prepare_directories_and_drop_privileges()
+    verify_cuda()
+
+    command = sys.argv[1:]
+    if command and command[0] == "--":
+        command = command[1:]
+    if command:
+        os.execvp(command[0], command)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -10,29 +10,27 @@ authors = [{name = "DeepSeek-AI"}]
 license = {file = "LICENSE-CODE"}
 urls = {homepage = "https://github.com/deepseek-ai/Janus"}
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.10,<3.13"
 dependencies = [
-    "torch>=2.0.1",
-    "transformers>=4.38.2",
-    "timm>=0.9.16",
-    "accelerate",
-    "sentencepiece",
-    "attrdict",
-    "einops",
+    "torch==2.9.1",
+    "torchvision==0.24.1",
+    "transformers==4.48.2",
+    "timm==0.9.16",
+    "accelerate==1.3.0",
+    "sentencepiece==0.2.1",
+    "attrdict3==2.0.2",
+    "einops==0.8.0",
 ]
 
 [project.optional-dependencies]
 gradio = [
-    "gradio==3.48.0",
-    "gradio-client==0.6.1",
-    "mdtex2html==1.3.0",
-    "pypinyin==0.50.0",
-    "tiktoken==0.5.2",
-    "tqdm==4.64.0",
-    "colorama==0.4.5",
-    "Pygments==2.12.0",
-    "markdown==3.4.1",
-    "SentencePiece==0.1.96"
+    "gradio==5.20.0",
+    "gradio-client==1.7.2",
+    "spaces==0.32.0",
+    "opencv-python-headless==4.11.0.86",
+    "matplotlib==3.10.0",
+    "pillow==11.1.0",
+    "pydantic==2.10.6",
 ]
 lint = [
     "isort",

requirements-gradio.txt

@@ -1,11 +1,3 @@
-# for gradio demo
-gradio==3.48.0
-gradio-client==0.6.1
-mdtex2html==1.3.0
-pypinyin==0.50.0
-tiktoken==0.5.2
-tqdm==4.64.0
-colorama==0.4.5
-Pygments==2.12.0
-markdown==3.4.1
-SentencePiece==0.1.96
+# Matches the Hugging Face Spaces SDK version declared in README.md.
+gradio==5.20.0
+gradio-client==1.7.2

requirements.txt

@@ -1,25 +1,24 @@
-torch>=2.0.1
-numpy>=1.21.2
+# GPU framework versions must match the pinned PyTorch CUDA base image.
+torch==2.9.1
+torchvision==0.24.1
+
+# The custom attention implementations in demo/modified_attn.py target the
+# Transformers 4.48 API. Do not upgrade this independently.
 transformers==4.48.2
-timm>=0.9.16
-accelerate
-sentencepiece
-attrdict
-einops
-opencv-python
-spaces
-matplotlib
-pillow>=11.1.0
-pydantic==2.10.6
+huggingface-hub==0.28.1
+tokenizers==0.21.0
+safetensors==0.5.2
+accelerate==1.3.0
 
-# for gradio demo
-# gradio==3.48.0
-# gradio-client==0.6.1
-# mdtex2html==1.3.0
-# pypinyin==0.50.0
-# tiktoken==0.5.2
-# tqdm==4.64.0
-# colorama==0.4.5
-# Pygments==2.12.0
-# markdown==3.4.1
-# SentencePiece==0.1.96
+# Model and visualization runtime.
+numpy==2.2.2
+timm==0.9.16
+sentencepiece==0.2.1
+attrdict3==2.0.2
+einops==0.8.0
+opencv-python-headless==4.11.0.86
+spaces==0.32.0
+matplotlib==3.10.0
+pillow==11.1.0
+pydantic==2.10.6
+psutil==5.9.8