Compile diffvg with CUDA

Dockerfile

@@ -35,8 +35,14 @@ RUN pip install --upgrade pip \
 # —— Build diffvg from a known-good commit ——————————
 RUN git clone https://github.com/BachiLi/diffvg.git /tmp/diffvg \
  && cd /tmp/diffvg \
+ && git checkout 7aa22d3f1ef72e22e5cfae3c8c50288f31bc9bcf \
  && git submodule update --init --recursive \
- && USE_CUDA=1 python setup.py install --user \
+ && mkdir build && cd build \
+ && cmake .. -DCMAKE_BUILD_TYPE=Release -DUSE_CUDA=ON \
+      -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \
+ && make -j$(nproc) \
+ && cd ../python && TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6;8.9" \
+      python setup.py install --user \
  && cd /app \
  && rm -rf /tmp/diffvg
 

README.md

@@ -57,9 +57,9 @@ cd Word-As-Image
 ```
 2. Create a new conda environment and install the libraries:
 ```bash
-conda create --name word python=3.8.15
+conda create --name word python=3.10
 conda activate word
-pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
+pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
 conda install -y numpy scikit-image
 conda install -y -c anaconda cmake
 conda install -y -c conda-forge ffmpeg
@@ -79,8 +79,14 @@ pip install transformers scipy ftfy accelerate
 ```bash
 git clone https://github.com/BachiLi/diffvg.git
 cd diffvg
+git checkout 7aa22d3f1ef72e22e5cfae3c8c50288f31bc9bcf
 git submodule update --init --recursive
-python setup.py install
+mkdir build && cd build
+cmake .. -DCMAKE_BUILD_TYPE=Release -DUSE_CUDA=ON \
+  -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc
+make -j$(nproc)
+cd ../python && TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6;8.9" \
+  python setup.py install
 ```
 
 5. Paste your HuggingFace [access token](https://huggingface.co/settings/tokens) for StableDiffusion in the TOKEN file.

code/main.py

@@ -55,9 +55,14 @@ def generate_word_image(cfg, device: torch.device):
     if isinstance(cfg, dict):
         cfg = edict(cfg)
 
-    # diffvg was compiled without GPU support in this environment. Always use
-    # the CPU for rendering to avoid runtime errors.
-    pydiffvg.set_use_gpu(False)
+    # Use the GPU for rendering when available.  If diffvg was compiled
+    # without CUDA support this call will fall back to CPU rendering.
+    use_gpu = torch.cuda.is_available()
+    try:
+        pydiffvg.set_use_gpu(use_gpu)
+    except Exception as err:  # diffvg not built with CUDA
+        print(f"Warning: cannot enable CUDA for diffvg ({err}). Falling back to CPU.")
+        pydiffvg.set_use_gpu(False)
 
     print(f"experiment directory: {cfg.experiment_dir}")