the editor part of the model editor API

Author: decahedron1

src/editor/mod.rs

@@ -255,7 +255,7 @@ impl Model {
 		Ok(())
 	}
 
-	pub fn into_session(self, options: SessionBuilder) -> Result<Session> {
+	pub fn into_session(self, mut options: SessionBuilder) -> Result<Session> {
 		let env = get_environment()?;
 		let mut session_ptr = ptr::null_mut();
 		ortsys![@editor:

src/session/builder/editable.rs

@@ -0,0 +1,65 @@
+use alloc::sync::Arc;
+use core::{
+	ops::Deref,
+	ptr::{self, NonNull}
+};
+
+use super::{PrepackedWeights, SessionBuilder};
+use crate::{AsPointer, Error, Result, editor::Model, ortsys, session::Session};
+
+pub struct EditableSession {
+	session: Session,
+	builder: SessionBuilder,
+	prepacked_weights: Option<PrepackedWeights>
+}
+
+impl EditableSession {
+	pub(crate) fn new(session: NonNull<ort_sys::OrtSession>, mut builder: SessionBuilder) -> Result<Self> {
+		// Prepacked weights are passed to `FinalizeModelEditorSession`; steal them from the builder so we can add them later.
+		let prepacked_weights = builder.prepacked_weights.take();
+		Ok(Self {
+			session: builder.commit_finalize(session)?,
+			builder,
+			prepacked_weights
+		})
+	}
+
+	pub fn apply_model(&mut self, model: &Model) -> Result<()> {
+		ortsys![@editor:
+			unsafe ApplyModelToModelEditorSession(
+				self.session.ptr_mut(),
+				model.ptr().cast_mut()
+			)?
+		];
+		Ok(())
+	}
+
+	pub fn into_session(mut self) -> Result<Session> {
+		ortsys![@editor:
+			unsafe FinalizeModelEditorSession(
+				self.session.ptr_mut(),
+				self.builder.ptr(),
+				self.prepacked_weights.as_ref().map(|p| p.ptr()).unwrap_or_else(ptr::null)
+			)?
+		];
+
+		if let Some(prepacked_weights) = self.prepacked_weights {
+			let Some(inner) = Arc::get_mut(&mut self.session.inner) else {
+				return Err(Error::new("Expected to have exclusive access to session inner"));
+			};
+
+			// add to extras so it outlives the session
+			inner._extras.push(Box::new(prepacked_weights));
+		}
+
+		Ok(self.session)
+	}
+}
+
+impl Deref for EditableSession {
+	type Target = Session;
+
+	fn deref(&self) -> &Self::Target {
+		&self.session
+	}
+}

src/session/builder/impl_commit.rs

@@ -15,7 +15,7 @@ use std::path::PathBuf;
 
 use smallvec::SmallVec;
 
-use super::SessionBuilder;
+use super::{EditableSession, SessionBuilder};
 #[cfg(feature = "std")]
 use crate::error::{Error, ErrorCode};
 use crate::{
@@ -176,7 +176,7 @@ impl SessionBuilder {
 		self.commit_finalize(unsafe { NonNull::new_unchecked(session_ptr) })
 	}
 
-	pub(crate) fn commit_finalize(mut self, ptr: NonNull<ort_sys::OrtSession>) -> Result<Session> {
+	pub(crate) fn commit_finalize(&mut self, ptr: NonNull<ort_sys::OrtSession>) -> Result<Session> {
 		let allocator = match &self.memory_info {
 			Some(info) => {
 				let mut allocator_ptr: *mut ort_sys::OrtAllocator = ptr::null_mut();
@@ -218,4 +218,47 @@ impl SessionBuilder {
 			outputs
 		})
 	}
+
+	#[cfg(feature = "std")]
+	#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
+	pub fn edit_from_file<P>(self, model_filepath: P) -> Result<EditableSession>
+	where
+		P: AsRef<Path>
+	{
+		let mut session_ptr: *mut ort_sys::OrtSession = ptr::null_mut();
+		let model_path = crate::util::path_to_os_char(model_filepath);
+
+		let env = get_environment()?;
+
+		ortsys![@editor:
+			unsafe CreateModelEditorSession(
+				env.ptr(),
+				model_path.as_ptr(),
+				self.session_options_ptr.as_ptr(),
+				&mut session_ptr
+			)?;
+			nonNull(session_ptr)
+		];
+
+		EditableSession::new(unsafe { NonNull::new_unchecked(session_ptr) }, self)
+	}
+
+	pub fn edit_from_memory(self, model_bytes: &[u8]) -> Result<EditableSession> {
+		let mut session_ptr: *mut ort_sys::OrtSession = ptr::null_mut();
+
+		let env = get_environment()?;
+
+		ortsys![@editor:
+			unsafe CreateModelEditorSessionFromArray(
+				env.ptr(),
+				model_bytes.as_ptr().cast(),
+				model_bytes.len() as _,
+				self.session_options_ptr.as_ptr(),
+				&mut session_ptr
+			)?;
+			nonNull(session_ptr)
+		];
+
+		EditableSession::new(unsafe { NonNull::new_unchecked(session_ptr) }, self)
+	}
 }

src/session/builder/mod.rs

@@ -8,11 +8,15 @@ use smallvec::SmallVec;
 
 use crate::{AsPointer, error::Result, logging::LoggerFunction, memory::MemoryInfo, operator::OperatorDomain, ortsys, util::with_cstr, value::DynValue};
 
+mod editable;
 mod impl_commit;
 mod impl_config_keys;
 mod impl_options;
 
-pub use self::impl_options::{GraphOptimizationLevel, PrepackedWeights};
+pub use self::{
+	editable::EditableSession,
+	impl_options::{GraphOptimizationLevel, PrepackedWeights}
+};
 
 /// Creates a session using the builder pattern.
 ///

src/session/mod.rs

@@ -32,7 +32,7 @@ use crate::{
 	memory::Allocator,
 	metadata::ModelMetadata,
 	ortsys,
-	util::{STACK_SESSION_INPUTS, STACK_SESSION_OUTPUTS, with_cstr_ptr_array},
+	util::{STACK_SESSION_INPUTS, STACK_SESSION_OUTPUTS, with_cstr, with_cstr_ptr_array},
 	value::{DynValue, Value, ValueType}
 };
 
@@ -545,6 +545,14 @@ impl Session {
 		ortsys![unsafe SetEpDynamicOptions(self.inner.session_ptr.as_ptr(), &key, &value, 1)?];
 		Ok(())
 	}
+
+	pub fn opset_for_domain(&self, domain: impl AsRef<str>) -> Result<u32> {
+		with_cstr(domain.as_ref().as_bytes(), &|domain| {
+			let mut opset = 0;
+			ortsys![@editor: unsafe SessionGetOpsetForDomain(self.inner.session_ptr.as_ptr(), domain.as_ptr(), &mut opset)?];
+			Ok(opset as u32)
+		})
+	}
 }
 
 /// Workload type, used to signal to execution providers whether to prioritize performance or efficiency.