Update app.py

app.py

@@ -10,6 +10,7 @@ import gradio as gr
 import sys
 from PIL import Image
 import importlib, site, sys
+import nvdiffrast.torch as dr
 
 # Re-discover all .pth/.egg-link files
 for sitedir in site.getsitepackages():
@@ -71,6 +72,8 @@ from omegaconf import OmegaConf
 
 DEVICE = "cuda"
 
+rasterizer = dr.RasterizeCudaContext()
+
 base_conf = OmegaConf.load(f'{env_paths.CODE_BASE}/configs/tracking.yaml')
 
 _mesh_file = env_paths.head_template
@@ -82,7 +85,8 @@ diff_renderer = NVDRenderer(
     image_size=base_conf.size, 
     obj_filename=_mesh_file,
     no_sh=False,
-    white_bg=True
+    white_bg=True,
+    glctx=rasterizer
 ).to(DEVICE)