fix: Handle log callback with None parameter in nanobind

- Use nb::handle instead of nb::callable to accept both callable and None
- Check is_none() explicitly before casting to callable
- Fixes segfault when callback is invoked after being cleared
- Store callback in unique_ptr to manage lifetime properly

Author: ONNX GenAI Assistant

src/python/python.cpp

@@ -323,15 +323,22 @@ void SetLogOptions(const nb::kwargs& dict) {
 }
 
 void SetLogCallback(std::optional<nb::callable> callback) {
-  static std::optional<nb::callable> log_callback;
-  log_callback = callback;
-
-  if (log_callback.has_value()) {
+  // Use a pointer to heap-allocated callable to avoid lifetime issues
+  static std::unique_ptr<nb::callable> log_callback_ptr;
+  
+  if (callback.has_value()) {
+    // Store the callback on the heap and keep it alive
+    log_callback_ptr = std::make_unique<nb::callable>(callback.value());
+    
     OgaPy::OgaCheckResult(OgaSetLogCallback([](const char* message, size_t length) {
-      nb::gil_scoped_acquire gil;
-      (*log_callback)(std::string_view(message, length));
+      if (log_callback_ptr) {
+        nb::gil_scoped_acquire gil;
+        (*log_callback_ptr)(std::string_view(message, length));
+      }
     }));
   } else {
+    // Clear the callback
+    log_callback_ptr.reset();
     OgaPy::OgaCheckResult(OgaSetLogCallback(nullptr));
   }
 }
@@ -862,7 +869,13 @@ NB_MODULE(onnxruntime_genai, m) {
       });
 
   m.def("set_log_options", &SetLogOptions);
-  m.def("set_log_callback", &SetLogCallback);
+  m.def("set_log_callback", [](nb::handle callback) {
+    if (callback.is_none()) {
+      SetLogCallback(std::nullopt);
+    } else {
+      SetLogCallback(nb::cast<nb::callable>(callback));
+    }
+  });
 
   m.def("is_cuda_available", []() { return USE_CUDA != 0; });
   m.def("is_dml_available", []() { return USE_DML != 0; });