From ce72ccd554a8be4d4d6b4169864659038697dfeb Mon Sep 17 00:00:00 2001
From: Jan Minar <rdancer@rdancer.org>
Date: Thu, 3 Oct 2024 20:50:41 +0100
Subject: [PATCH] fix: No available kernel. Aborting execution.

---
 __init__.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index 46e879f..3e54efb 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,6 +1,7 @@
 import torch
 from PIL import Image
 import numpy as np
+import os
 
 try:
     from app import process_image
@@ -30,7 +31,15 @@ def pil2tensor(image: Image.Image) -> torch.Tensor:
 
 class F2S2GenerateMask:
     def __init__(self):
-        pass
+        self._enable_cudnn_sdpa()
+
+    def _enable_cudnn_sdpa(self):
+        # Check if the environment variable is already set
+        if os.getenv('TORCH_CUDNN_SDPA_ENABLED') != '1':
+            os.environ['TORCH_CUDNN_SDPA_ENABLED'] = '1'
+            print("CuDNN SDPA enabled.")
+        else:
+            print("CuDNN SDPA was already enabled.")
 
     @classmethod
     def INPUT_TYPES(cls):