comfy-builder-core 0.0.2

use candle_core::shape::ShapeWithOneHole;
use candle_core::{Device, Tensor as CandleTensor, WithDType};
use numpy::{Element, PyArray, PyArrayDyn, PyArrayMethods, PyUntypedArrayMethods};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::PyAnyMethods;
use pyo3::{Bound, FromPyObject, IntoPyObject, PyAny, PyErr, PyResult, Python};
use std::marker::PhantomData;
use std::ops::Deref;

#[derive(Clone, Debug)]
pub struct Mask<T: Element + WithDType = f32>(Tensor<T>);

impl<T: Element + WithDType> Deref for Mask<T> {
    type Target = CandleTensor;

    fn deref(&self) -> &Self::Target {
        &self.0.tensor
    }
}

#[derive(Clone, Debug)]
pub struct Tensor<T = f32>
where
    T: Element + WithDType,
{
    tensor: CandleTensor,
    _marker: PhantomData<T>,
}

impl<T> Tensor<T>
where
    T: Element + WithDType,
{
    pub fn new<'py>(py_any: &Bound<'py, PyAny>, device: &Device) -> Self {
        let tensor = Self::torch_to_candle(py_any, device);

        Self {
            tensor,
            _marker: PhantomData,
        }
    }

    fn torch_to_candle<'py>(torch_tensor: &Bound<'py, PyAny>, device: &Device) -> CandleTensor {
        let mut np = torch_tensor.call_method0("numpy").unwrap();

        let mut arr = np.downcast::<PyArrayDyn<T>>().unwrap();

        if !arr.is_contiguous() {
            np = np.call_method0("copy").unwrap();
            arr = np.downcast::<PyArrayDyn<T>>().unwrap();
        }

        let shape = arr.shape().to_vec();
        let data = arr.to_vec().unwrap();

        CandleTensor::from_vec(data, shape, device).unwrap()
    }

    /// The dimension size for this tensor on each axis.
    pub fn dims(&self) -> &[usize] {
        self.tensor.dims()
    }

    pub fn from_tensor(tensor: CandleTensor) -> Self {
        Self {
            tensor,
            _marker: PhantomData,
        }
    }

    pub fn into_tensor(self) -> CandleTensor {
        self.tensor
    }
}

impl<'py, T> IntoPyObject<'py> for Tensor<T>
where
    T: Element + WithDType,
{
    type Target = PyAny;
    type Output = Bound<'py, Self::Target>;
    type Error = PyErr;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        let tensor = self.into_tensor();
        let shape = tensor.dims();

        let data: Vec<T> = tensor
            .flatten_all()
            .map_err(|e| PyErr::new::<PyRuntimeError, _>(e.to_string()))?
            .to_vec1::<_>()
            .map_err(|e| PyErr::new::<PyRuntimeError, _>(e.to_string()))?;

        let array = PyArray::from_iter(py, data)
            .reshape(shape)
            .map_err(|e| PyErr::new::<PyRuntimeError, _>(e.to_string()))?;

        let torch = py.import("torch")?;
        let tensor = torch.getattr("tensor")?.call1((array,))?;

        Ok(tensor)
    }
}

impl<'py, T> IntoPyObject<'py> for Mask<T>
where
    T: Element + WithDType,
{
    type Target = PyAny;
    type Output = Bound<'py, Self::Target>;
    type Error = PyErr;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        self.0.into_pyobject(py)
    }
}

impl<T: Element + WithDType> Tensor<T> {
    pub fn from_raw<U: ShapeWithOneHole>(
        data: Vec<T>,
        shape: U,
        device: &Device,
    ) -> candle_core::Result<Tensor> {
        Ok(Tensor::from_tensor(CandleTensor::from_vec(
            data, shape, device,
        )?))
    }
}

impl<T: Element + WithDType, S: ShapeWithOneHole> TryFrom<(Vec<T>, S, &Device)> for Tensor {
    type Error = candle_core::Error;

    fn try_from(value: (Vec<T>, S, &Device)) -> Result<Self, Self::Error> {
        Tensor::from_raw(value.0, value.1, value.2)
    }
}

impl<T: Element + WithDType, S: ShapeWithOneHole> TryFrom<(Vec<T>, S, &Device)> for Mask {
    type Error = candle_core::Error;

    fn try_from(value: (Vec<T>, S, &Device)) -> Result<Self, Self::Error> {
        Ok(Mask(Tensor::from_raw(value.0, value.1, value.2)?))
    }
}

impl<T: Element + WithDType> Deref for Tensor<T> {
    type Target = CandleTensor;

    fn deref(&self) -> &Self::Target {
        &self.tensor
    }
}

impl<'py, T> FromPyObject<'py> for Tensor<T>
where
    T: Element + WithDType,
{
    fn extract_bound(object: &Bound<'py, PyAny>) -> PyResult<Self> {
        object
            .extract::<Bound<'py, PyAny>>()
            .map(|value| Tensor::new(&value, &Device::Cpu))
    }
}

impl<'py, T: Element + WithDType> FromPyObject<'py> for Mask<T> {
    fn extract_bound(object: &Bound<'py, PyAny>) -> PyResult<Self> {
        object
            .extract::<Bound<'py, PyAny>>()
            .map(|value| Tensor::new(&value, &Device::Cpu).into())
    }
}

impl<T: Element + WithDType> From<Tensor<T>> for Mask<T> {
    fn from(tensor: Tensor<T>) -> Self {
        Mask(tensor)
    }
}