diff --git a/nvim/.config/nvim/colors/yolokai.vim b/nvim/.config/nvim/colors/yolokai.vim
index b370e20..561dcc3 100644
--- a/nvim/.config/nvim/colors/yolokai.vim
+++ b/nvim/.config/nvim/colors/yolokai.vim
@@ -590,6 +590,9 @@ local highlight_groups = {
   gitrebaseReword        = 'gitrebasePick'        ,
   gitrebaseSquash        = 'Macro'                ,
   gitrebaseSummary       = 'Normal'               ,
+  -- Needed for treesitter diff
+  ['@text.diff.add']     = 'DiffAdd'              ,
+  ['@text.diff.delete']  = 'DiffDelete'           ,
 
   --[[ 4.3.35 Help ]]
   helpHeader        = 'Label'     ,