Skip to content

Run ditto-talkinghead in a Docker with 5090 #88

Description

@cuso03

HI everyone, the last weeks I notice about this model so first i tried in the collab and then on mi Desktop with pytorch and everything its okay, but when i want to make a docker to integrate this model in a fast api project i have problems with tensort so i must to make an environment (in the dockerfile) which have python 3.10 , pythorch 2.5 and the tensort packages with 8.6.1. After that when i tried to run the image the terminal returns this log:

root@gpu-station:/app/ditto-talkinghead#` python inference.py \
    --data_root "../checkpoints/ditto_pytorch" \
    --cfg_pkl "../checkpoints/ditto_cfg/v0.4_hubert_cfg_pytorch.pkl" \
    --audio_path "./example/audio.wav" \
    --source_path "./example/image.png" \
    --output_path "./tmp/result_docker.mp4" 
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1776706821.459688     157 gl_context_egl.cc:85] Successfully initialized EGL. Major : 1 Minor: 5
I0000 00:00:1776706821.494588     310 gl_context.cc:357] GL version: 3.2 (OpenGL ES 3.2 NVIDIA 580.126.09), renderer: NVIDIA GeForce RTX 5090/PCIe/SSE2
W0000 00:00:1776706821.496429     157 face_landmarker_graph.cc:174] Sets FaceBlendshapesGraph acceleration to xnnpack by default.
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.
W0000 00:00:1776706821.507227     315 inference_feedback_manager.cc:114] Feedback manager requires a model with a single signature inference. Disabling support for feedback tensors.
W0000 00:00:1776706821.522258     320 inference_feedback_manager.cc:114] Feedback manager requires a model with a single signature inference. Disabling support for feedback tensors.
/app/ditto-talkinghead/core/models/modules/appearance_feature_extractor.py:72: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  self.load_state_dict(torch.load(ckpt_path, map_location=lambda storage, loc: storage))
/usr/local/lib/python3.10/dist-packages/torch/cuda/__init__.py:235: UserWarning: 
NVIDIA GeForce RTX 5090 with CUDA capability sm_120 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_50 sm_60 sm_70 sm_75 sm_80 sm_86 sm_90.
If you want to use the NVIDIA GeForce RTX 5090 GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

  warnings.warn(
/app/ditto-talkinghead/core/models/modules/motion_extractor.py:23: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  self.load_state_dict(torch.load(ckpt_path, map_location=lambda storage, loc: storage))
/app/ditto-talkinghead/core/models/modules/LMDM.py:84: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  checkpoint = torch.load(ckpt_path, map_location='cpu')
/app/ditto-talkinghead/core/models/modules/stitching_network.py:44: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
/app/ditto-talkinghead/core/models/modules/warping_network.py:85: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  self.load_state_dict(torch.load(ckpt_path, map_location=lambda storage, loc: storage))
/app/ditto-talkinghead/core/models/modules/spade_generator.py:85: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  self.load_state_dict(torch.load(ckpt_path, map_location=lambda storage, loc: storage))
==================== setup kwargs ====================
max_size <class 'int'> 1920
template_n_frames <class 'int'> -1
crop_scale <class 'float'> 2.3
crop_vx_ratio <class 'int'> 0
crop_vy_ratio <class 'float'> -0.125
crop_flag_do_rot <class 'bool'> True
smo_k_s <class 'int'> 13
emo <class 'numpy.ndarray'> (600, 8)
eye_f0_mode <class 'bool'> False
ch_info <class 'dict'>
overlap_v2 <class 'int'> 10
fix_kp_cond <class 'int'> 1
fix_kp_cond_dim <class 'list'> [0, 202]
sampling_timesteps <class 'int'> 50
online_mode <class 'bool'> False
v_min_max_for_clip <class 'numpy.ndarray'> (4, 265)
smo_k_d <class 'int'> 3
N_d <class 'int'> -1
use_d_keys <class 'NoneType'> None
relative_d <class 'bool'> True
drive_eye <class 'NoneType'> None
delta_eye_arr <class 'numpy.ndarray'> (15, 63)
delta_eye_open_n <class 'int'> 0
fade_type <class 'str'> d0
fade_out_keys <class 'list'> ['exp']
flag_stitching <class 'bool'> True
overall_ctrl_info <class 'dict'> {'delta_pitch': 2}
==================================================

2026-04-20 17:41:32.857640818 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {800,10} does not match actual shape of {512,10} for output 500
2026-04-20 17:41:34.633621354 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {800,4} does not match actual shape of {512,4} for output 497
2026-04-20 17:41:34.634445962 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {800,1} does not match actual shape of {512,1} for output 494
2026-04-20 17:41:34.641886595 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {3200,10} does not match actual shape of {2048,10} for output 477
2026-04-20 17:41:34.642880974 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {3200,4} does not match actual shape of {2048,4} for output 474
2026-04-20 17:41:34.643780233 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {3200,1} does not match actual shape of {2048,1} for output 471
2026-04-20 17:41:34.657608924 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {12800,10} does not match actual shape of {8192,10} for output 454
2026-04-20 17:41:34.658979102 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {12800,4} does not match actual shape of {8192,4} for output 451
2026-04-20 17:41:34.660415163 [W:onnxruntime:, execution_frame.cc:870 VerifyOutputSizes] Expected shape from model of {12800,1} does not match actual shape of {8192,1} for output 448
W0000 00:00:1776706900.913797     319 landmark_projection_calculator.cc:186] Using NORM_RECT without IMAGE_DIMENSIONS is only supported for the square ROI. Provide IMAGE_DIMENSIONS or use PROJECTION_MATRIX.
Traceback (most recent call last):
  File "/app/ditto-talkinghead/inference.py", line 89, in <module>
    run(SDK, audio_path, source_path, output_path)
  File "/app/ditto-talkinghead/inference.py", line 34, in run
    SDK.setup(source_path, output_path, **setup_kwargs)
  File "/app/ditto-talkinghead/stream_pipeline_offline.py", line 148, in setup
    source_info = self.avatar_registrar(
  File "/app/ditto-talkinghead/core/atomic_components/avatar_registrar.py", line 101, in __call__
    return self.register(*args, **kwargs)
  File "/app/ditto-talkinghead/core/atomic_components/avatar_registrar.py", line 86, in register
    info = self.source2info(rgb, last_lmk, **kwargs)
  File "/app/ditto-talkinghead/core/atomic_components/source2info.py", line 144, in __call__
    kp_info = self._get_kp_info(rgb_256_bchw)
  File "/app/ditto-talkinghead/core/atomic_components/source2info.py", line 113, in _get_kp_info
    kp_info = self.motion_extractor(img)
  File "/app/ditto-talkinghead/core/models/motion_extractor.py", line 40, in __call__
    pred = self.model(torch.from_numpy(image).to(self.device))
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/app/ditto-talkinghead/core/models/modules/motion_extractor.py", line 19, in forward
    out = self.detector(x)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/app/ditto-talkinghead/core/models/modules/convnextv2.py", line 120, in forward
    x = self.forward_features(x)
  File "/app/ditto-talkinghead/core/models/modules/convnextv2.py", line 115, in forward_features
    x = self.downsample_layers[i](x)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/container.py", line 250, in forward
    input = module(input)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py", line 554, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py", line 549, in _conv_forward
    return F.conv2d(
RuntimeError: CUDA error: no kernel image is available for execution on the device
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Since I don't have much experience i ask to claude and tells me that is because mi graphic card works with the cuda 12.8 Blackwell image, but this means that i should work with a modern stack with a higher version of pytorch and other dependencies. Im very confused and i dont know what i can to do. All type off help is welcome.

this is mi current Dockerfile:

# =============================================================================
# ditto-talkinghead — Dockerfile
# Base: CUDA 12.1.1 + cuDNN 8.x (requerido por TensorRT 8.6.1)
# Python: 3.10 | PyTorch: 2.5.1+cu121
# =============================================================================

FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

# -- Variables de entorno globales --------------------------------------------
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV PIP_NO_CACHE_DIR=1
ENV PIP_DISABLE_PIP_VERSION_CHECK=1

# cuDNN 8 ya está en la imagen base; apuntamos las librerías explícitamente
# para que TensorRT las encuentre sin ambigüedad
ENV CUDNN_PATH=/usr/lib/x86_64-linux-gnu
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${CUDNN_PATH}:${LD_LIBRARY_PATH}

# -- Dependencias del sistema -------------------------------------------------
# Ubuntu 22.04 ya trae Python 3.10; solo instalamos pip y libs necesarias.
# NO usamos add-apt-repository para evitar dependencia de red a PPAs externos.
RUN apt-get update && apt-get install -y --no-install-recommends \
    # Utilidades base
    curl \
    git \
    wget \
    ca-certificates \
    # Python 3.10 (viene con Ubuntu 22.04)
    python3.10 \
    python3.10-dev \
    python3-pip \
    # Video / audio
    ffmpeg \
    libsndfile1 \
    libsndfile1-dev \
    # OpenCV headless requiere estas libs de sistema
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender1 \
    # Build tools (insightface compila C extensions)
    build-essential \
    cmake \
    ninja-build \
    # Misc
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# Crear alias python -> python3.10
RUN ln -sf /usr/bin/python3.10 /usr/bin/python \
    && ln -sf /usr/bin/python3.10 /usr/bin/python3 \
    && python -m pip install --upgrade pip setuptools wheel

# -- PyTorch (instalado primero, aislado del resto) ---------------------------
# Trío oficial torch+torchvision+torchaudio para cu121.
# torchvision 0.20.1 es el par oficial de torch 2.5.1 (reemplaza el dev+cu124).
RUN pip install \
    torch==2.5.1+cu121 \
    torchvision==0.20.1+cu121 \
    torchaudio==2.5.1+cu121 \
    --index-url https://download.pytorch.org/whl/cu121

# -- TensorRT 8.6.1 -----------------------------------------------------------
# Instalado antes del resto para que las dependencias lo vean como base.
# Usa el cuDNN 8.x de la imagen base; NO instalar nvidia-cudnn-cu12 por pip
# (genera conflicto de versiones con el cuDNN del sistema).
COPY wheels/ /tmp/wheels/

# 1. Instalamos las librerías pesadas y bindings que SÍ son binarios completos (.whl)
RUN pip install --no-cache-dir /tmp/wheels/tensorrt_libs-*.whl
RUN pip install --no-cache-dir /tmp/wheels/tensorrt_bindings-*.whl

# 2. El paquete 'tensorrt' (el .tar.gz) falla porque quiere internet.
# Lo descomprimimos nosotros y movemos la carpeta del código a donde Python la busca.
RUN mkdir -p /tmp/trt_extracted && \
    tar -xzf /tmp/wheels/tensorrt-8.6.1.tar.gz -C /tmp/trt_extracted && \
    # Movemos solo la carpeta 'tensorrt' al directorio de paquetes de Python
    cp -r /tmp/trt_extracted/tensorrt-8.6.1/tensorrt /usr/local/lib/python3.10/dist-packages/ && \
    rm -rf /tmp/trt_extracted


# -- Resto de dependencias pip ------------------------------------------------
COPY requirements.txt /tmp/requirements.txt
RUN pip install -r /tmp/requirements.txt

# -- Código de la aplicación --------------------------------------------------
WORKDIR /app
COPY . /app

# Los checkpoints/modelos se montan como volumen externo (no se incluyen
# en la imagen para mantenerla portable y no depender de red en el build).
VOLUME ["/app/checkpoints", "/app/models"]

# Puerto por si se expone algún servicio (ajustar según necesidad)
EXPOSE 8080

CMD ["bash"]

The Dockerfile is made it with the "ofline build" where i download the model, some dependencies locally in a secundary computer then copy in the Dockerfile , make the build , save the image in a tar and then load in the workstation. Because in the workstation where i work i have problems with the dns.

SOME EXTRA DETAILS :

| NVIDIA-SMI 580.126.09             Driver Version: 580.126.09     CUDA Version: 13.0     |
+-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 5090        Off |   00000000:01:00.0  On |                  N/A |
|  0%   41C    P8             15W /  600W |     383MiB /  32607MiB |      1%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions