comfy_builder_core/
tensor.rs

1use candle_core::shape::ShapeWithOneHole;
2use candle_core::{Device, Tensor as CandleTensor, WithDType};
3use numpy::{Element, PyArray, PyArrayDyn, PyArrayMethods, PyUntypedArrayMethods};
4use pyo3::exceptions::PyRuntimeError;
5use pyo3::prelude::PyAnyMethods;
6use pyo3::{Bound, FromPyObject, IntoPyObject, PyAny, PyErr, PyResult, Python};
7use std::marker::PhantomData;
8use std::ops::Deref;
9
10#[derive(Clone, Debug)]
11pub struct Mask<T: Element + WithDType = f32>(Tensor<T>);
12
13impl<T: Element + WithDType> Deref for Mask<T> {
14    type Target = CandleTensor;
15
16    fn deref(&self) -> &Self::Target {
17        &self.0.tensor
18    }
19}
20
21#[derive(Clone, Debug)]
22pub struct Tensor<T = f32>
23where
24    T: Element + WithDType,
25{
26    tensor: CandleTensor,
27    _marker: PhantomData<T>,
28}
29
30impl<T> Tensor<T>
31where
32    T: Element + WithDType,
33{
34    pub fn new<'py>(py_any: &Bound<'py, PyAny>, device: &Device) -> Self {
35        let tensor = Self::torch_to_candle(py_any, device);
36
37        Self {
38            tensor,
39            _marker: PhantomData,
40        }
41    }
42
43    fn torch_to_candle<'py>(torch_tensor: &Bound<'py, PyAny>, device: &Device) -> CandleTensor {
44        let mut np = torch_tensor.call_method0("numpy").unwrap();
45
46        let mut arr = np.downcast::<PyArrayDyn<T>>().unwrap();
47
48        if !arr.is_contiguous() {
49            np = np.call_method0("copy").unwrap();
50            arr = np.downcast::<PyArrayDyn<T>>().unwrap();
51        }
52
53        let shape = arr.shape().to_vec();
54        let data = arr.to_vec().unwrap();
55
56        CandleTensor::from_vec(data, shape, device).unwrap()
57    }
58
59    /// The dimension size for this tensor on each axis.
60    pub fn dims(&self) -> &[usize] {
61        self.tensor.dims()
62    }
63
64    pub fn from_tensor(tensor: CandleTensor) -> Self {
65        Self {
66            tensor,
67            _marker: PhantomData,
68        }
69    }
70
71    pub fn into_tensor(self) -> CandleTensor {
72        self.tensor
73    }
74}
75
76impl<'py, T> IntoPyObject<'py> for Tensor<T>
77where
78    T: Element + WithDType,
79{
80    type Target = PyAny;
81    type Output = Bound<'py, Self::Target>;
82    type Error = PyErr;
83
84    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
85        let tensor = self.into_tensor();
86        let shape = tensor.dims();
87
88        let data: Vec<T> = tensor
89            .flatten_all()
90            .map_err(|e| PyErr::new::<PyRuntimeError, _>(e.to_string()))?
91            .to_vec1::<_>()
92            .map_err(|e| PyErr::new::<PyRuntimeError, _>(e.to_string()))?;
93
94        let array = PyArray::from_iter(py, data)
95            .reshape(shape)
96            .map_err(|e| PyErr::new::<PyRuntimeError, _>(e.to_string()))?;
97
98        let torch = py.import("torch")?;
99        let tensor = torch.getattr("tensor")?.call1((array,))?;
100
101        Ok(tensor)
102    }
103}
104
105impl<'py, T> IntoPyObject<'py> for Mask<T>
106where
107    T: Element + WithDType,
108{
109    type Target = PyAny;
110    type Output = Bound<'py, Self::Target>;
111    type Error = PyErr;
112
113    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
114        self.0.into_pyobject(py)
115    }
116}
117
118impl<T: Element + WithDType> Tensor<T> {
119    pub fn from_raw<U: ShapeWithOneHole>(
120        data: Vec<T>,
121        shape: U,
122        device: &Device,
123    ) -> candle_core::Result<Tensor> {
124        Ok(Tensor::from_tensor(CandleTensor::from_vec(
125            data, shape, device,
126        )?))
127    }
128}
129
130impl<T: Element + WithDType, S: ShapeWithOneHole> TryFrom<(Vec<T>, S, &Device)> for Tensor {
131    type Error = candle_core::Error;
132
133    fn try_from(value: (Vec<T>, S, &Device)) -> Result<Self, Self::Error> {
134        Tensor::from_raw(value.0, value.1, value.2)
135    }
136}
137
138impl<T: Element + WithDType, S: ShapeWithOneHole> TryFrom<(Vec<T>, S, &Device)> for Mask {
139    type Error = candle_core::Error;
140
141    fn try_from(value: (Vec<T>, S, &Device)) -> Result<Self, Self::Error> {
142        Ok(Mask(Tensor::from_raw(value.0, value.1, value.2)?))
143    }
144}
145
146impl<T: Element + WithDType> Deref for Tensor<T> {
147    type Target = CandleTensor;
148
149    fn deref(&self) -> &Self::Target {
150        &self.tensor
151    }
152}
153
154impl<'py, T> FromPyObject<'py> for Tensor<T>
155where
156    T: Element + WithDType,
157{
158    fn extract_bound(object: &Bound<'py, PyAny>) -> PyResult<Self> {
159        object
160            .extract::<Bound<'py, PyAny>>()
161            .map(|value| Tensor::new(&value, &Device::Cpu))
162    }
163}
164
165impl<'py, T: Element + WithDType> FromPyObject<'py> for Mask<T> {
166    fn extract_bound(object: &Bound<'py, PyAny>) -> PyResult<Self> {
167        object
168            .extract::<Bound<'py, PyAny>>()
169            .map(|value| Tensor::new(&value, &Device::Cpu).into())
170    }
171}
172
173impl<T: Element + WithDType> From<Tensor<T>> for Mask<T> {
174    fn from(tensor: Tensor<T>) -> Self {
175        Mask(tensor)
176    }
177}
comfy_builder_core/tensor.rs

comfy_builder_core/
tensor.rs
