kornia.core#

class kornia.core.TensorWrapper(data)[source]#

Wrapper around PyTorch tensors that tracks attribute and function usage.

This class provides a transparent wrapper around PyTorch tensors while tracking which attributes and functions are accessed. Useful for debugging and understanding tensor usage patterns.

Convention:

An attribute the wrapper does not define is looked up on the wrapped tensor, and a tensor value is wrapped in the wrapper’s class. The wrapper’s own names are never forwarded: its slots, data, __dict__ and the copy and pickle hooks. So copy.copy, copy.deepcopy and pickle return the wrapper’s class, and torch.compile can trace a function that builds a wrapper. The array and DLPack hooks are forwarded, so numpy.asarray(w) and torch.from_dlpack(w) convert the wrapped tensor.
A torch function called on a wrapper (torch.clone(w)), a tensor method reached through the wrapper (w.clone()), a tensor attribute (w.T) and every arithmetic, bitwise and comparison operator (+ - * / // % ** @ & | ^ << >>, == != < <= > >=, unary -, +, abs and ~) compute the wrapped tensors’ result and wrap each tensor in it, or in a list or tuple of them, by one rule. A torch.return_types result (torch.sort, torch.max with dim) is returned as torch returns it, with plain tensors, since the rule would also wrap its indices. The class of an operand with a coordinate width (Vector2, Vector3) is taken, whichever side it is on, when the result ends in that width and keeps at least that operand’s axes. Otherwise the class of the first operand without a width (Scalar, TensorWrapper) is taken: the left operand’s (w + x), or the right operand’s when the left one is not a wrapper (2 / w, t + w). A comparison whose right operand is an instance of a subclass of the left operand’s class takes the subclass, because Python tries the subclass’s reflected comparison first. When no operand qualifies, the result is a plain tensor. So Scalar * Vector3 is a Vector3, and torch.linalg.norm(v, dim=-1) is a plain tensor, also when it happens to end in 3. The rule reads shapes only: a result of the same shape stays in the class, as the transpose of a (3, 3) Vector3 does.
Indexing returns the indexed tensor in the wrapper’s class. For a subclass with a coordinate width, it does so only when the key leaves the coordinate axis last, at most sliced, and the result ends in the width (v[0], v[mask]); otherwise it returns a plain tensor (v[..., 0], v[:, 0]).
An in-place operator (+=, -=, *=, /=, //=, %=, **=, &=, |=, ^=, <<=, >>=) updates the wrapped tensor in place and returns the same wrapper, so an alias sees the change. The wrapper does not copy the tensor it is built from, so the update also changes that tensor, and it raises where the tensor’s in-place operator raises: on a leaf that requires grad, or when the result would change the dtype or the shape. w @= x rebinds w to w @ x, as @= does for a tensor.
bool, int, float, complex, operator.index and len return Python values.

Variables:

_data – The underlying PyTorch tensor.
used_attrs – Set of attribute names that have been accessed.
used_calls – Set of functions that have been called.

property data: Tensor#: Access the underlying tensor.

unwrap()[source]#

Return the underlying PyTorch tensor.

Return type:: Tensor

used_attrs#

used_calls#

kornia.core.utils.batched_forward(model, data, device, batch_size=128, **kwargs)[source]#

Run the forward in micro-batches.

When the just model.forward(data) does not fit into device memory, e.g. on laptop GPU. In the end, it transfers the output to the device of the input data tensor. E.g. running HardNet on 8000x1x32x32 tensor.

Removed from kornia.utils.memory in 0.8.3 and restored here as public API.

Parameters:

model (Module) – Any torch model, which outputs a single tensor as an output.
data (Tensor) – Input data of Bx(Any) shape.
device (device) – which device should we run on.
batch_size (int, optional) – “micro-batch” size. Default: 128
**kwargs (Any) – any other arguments, which accepts model.

Return type:

Tensor

Returns:

output of the model.

Example

>>> import torch
>>> from kornia.core.utils import batched_forward
>>> model = torch.nn.Identity()
>>> x = torch.rand(300, 2)
>>> out = batched_forward(model, x, torch.device("cpu"), batch_size=128)
>>> bool(torch.allclose(out, x))
True

Device and autocast helpers#

kornia.core.utils.get_cuda_device_if_available(index=0)[source]#

Try to get cuda device, if fail, return cpu.

Parameters:: index (int, optional) – cuda device index Default: 0
Return type:: device
Returns:: torch.device

kornia.core.utils.get_mps_device_if_available()[source]#

Try to get mps device, if fail, return cpu.

Return type:: device
Returns:: torch.device

kornia.core.utils.get_cuda_or_mps_device_if_available()[source]#

Check OS and platform and run get_cuda_device_if_available or get_mps_device_if_available.

Return type:: device
Returns:: torch.device

kornia.core.utils.is_autocast_enabled(both=True)[source]#

Check if torch autocast is enabled.

Parameters:: both (bool, optional) – if True, consider the autocast regions of the CPU, CUDA, MPS and XPU device types. Default: True
Return type:: bool
Returns:: If both is True, whether autocast is enabled for any of the CPU, CUDA, MPS or XPU device types, on every supported torch version. If both is False, torch.is_autocast_enabled() without a device type, whose device type depends on the torch version: it never reports CPU autocast, and on torch 2.5.1 it does not report MPS autocast either.

Weights#

Helpers for fetching and reading pretrained weights. load_state_dict_from_url() is the one models load a torch.load-able checkpoint with; download_hf_file() (or download_file_from_url()) and load_safetensors() are the two halves of the same job for a .safetensors checkpoint; check_safetensors() is the validate= callable that lets the download half reject a truncated cache entry.

kornia.core.hf_url(repo, filename)[source]#

Return the HuggingFace URL for a file in a model repo.

Both arguments are percent-encoded, / excepted, because they are path segments rather than URL syntax: unencoded, the # of "a#b.pth" would start a fragment and fetch the file a, and a space would make the URL invalid. Names made of letters, digits, -, _, . and ~ are unchanged. Pass the plain name: one that is already percent-encoded, such as "a%20b.pth", is encoded again.

Parameters:

repo (str) – repository name under the kornia HF org (e.g. "hardnet"), or a full owner/name repository id for a repo owned by anyone else (e.g. "google/siglip2-base-patch16-224"). The two are told apart by the /, which a repository name cannot contain.
filename (str) – path of the file in that repo (e.g. "HardNetPP.pth"); a / in it reaches into a subdirectory.

Return type:

str

Returns:

A resolve/main URL that can be passed directly to load_state_dict_from_url() or download_file_from_url().

Example

>>> hf_url("hardnet", "HardNetPP.pth")
'https://huggingface.co/kornia/hardnet/resolve/main/HardNetPP.pth'
>>> hf_url("google/siglip2-base-patch16-224", "model.safetensors")
'https://huggingface.co/google/siglip2-base-patch16-224/resolve/main/model.safetensors'

kornia.core.load_state_dict_from_url(url, *, timeout=None, **kwargs)[source]#

Load a state dict from a URL, trying fallback URLs on failure.

Replacement for torch.hub.load_state_dict_from_url() that also accepts an ordered list of URLs. Each URL is tried in turn; a warnings message is emitted for every failed attempt before the next source is tried. It deliberately differs from the torch function in two ways, both described below: the weights_only default and where progress is reported.

The checkpoint is loaded with weights_only=True unless the caller passes weights_only=False, whereas the torch function defaults to False on every torch version kornia supports; weights_only=True is torch.load’s own default since torch 2.6. torch.load then unpickles only tensors, primitive types and plain containers, and refuses a pickled callable instead of running it. A checkpoint that stores any other type fails with a RuntimeError chained to the pickle.UnpicklingError that names the type. Allowlist the type for the call with torch.serialization.safe_globals([...]), or pass weights_only=False, but only for a file you trust, because unpickling it can run arbitrary code.

On older torch, weights_only=True narrows what a checkpoint can do but does not guarantee that it runs no code. PyTorch’s advisories report checkpoints crafted to run code despite it before torch 2.6 (GHSA-53q9-r3pm-6pq6) and to corrupt memory, potentially running code, before torch 2.10 (GHSA-63cw-57p8-fm3p). kornia supports torch 2.5.1 and later, so load a checkpoint from a source you do not trust only on torch 2.10 or later, which fixes both.

Progress reporting is written to sys.stderr. This is the second deliberate deviation from the torch function, which since torch 2.x writes its Downloading: "<url>" to <path> line to sys.stdout (the accompanying progress bar already goes to stderr). Status output on stdout corrupts any caller that treats stdout as data – most visibly doctests, where the line is captured as unexpected example output and fails an example that downloads on a cold cache.

That line is reached only when the file is absent from the cache, so this function fetches a missing file itself – announcing it on stderr – and leaves torch with nothing to report. Nothing process-global is touched: redirecting sys.stdout around the call would divert unrelated threads’ output for the whole transfer, and concurrent calls restoring out of order would leave stdout permanently pointing at stderr.

When multiple URLs are given and file_name is not already in kwargs, the basename of the first URL is used as the local cache filename for all attempts. This guarantees that:

a file successfully downloaded from the primary source is found on the next call without re-downloading;
hash validation (check_hash=True) uses the filename — and therefore the hash embedded in it — of the primary URL consistently across all fallback attempts.

Because that one path is shared, a failed attempt moves the cache entry aside before the next URL is tried. Otherwise a single bad write would be handed straight back to torch by every remaining source – and by every later process – making the fallback URLs unreachable. If no source ends up transferring the file, the source the entry was moved aside for is tried once more, this time against an empty path, so a poisoned entry is repaired inside the call rather than after one guaranteed spurious failure – which matters most for the single-source checkpoints, where there is no fallback to do it. Should the call still fail, the entry is put back, so a failure the cache could not have fixed never costs the caller a checkpoint it already had. A path is refetched at most once per process this way; see _discard_cache_entry() and _settle_quarantine().

Only an entry that predates the call is ever quarantined. Bytes a source transferred itself and then could not load are not something the caller had, so there is nothing to make reversible: they are deleted outright while a later source remains, which is what hands that source an empty path to fetch into (see _drop_failed_download()). After the last source they are left where they are – deleting them would make every later call in the process transfer the file again, unbounded and up to 2.4 GB a time, to fail in the same way – but no discard is spent on them either, so the next call can still move them aside and reach the sources behind them.

Each URL is attempted up to _MAX_ATTEMPTS times, with exponential backoff, when it fails with a transient network condition such as an HTTP 429. Unauthenticated CI matrices trip the rate limits of huggingface.co and github.com routinely, and a retry is far cheaper than a failed job. A response that names its own delay in Retry-After – or, on a rate-limit response, X-RateLimit-Reset – is honoured instead of the guess, clamped to _MAX_BACKOFF_SECONDS, and the call as a whole never waits longer than _MAX_CALL_SLEEP_SECONDS. A server that stalls – accepts the connection and then sends nothing – fails the attempt after timeout seconds and counts as transient too.

The arguments are checked before the cache is consulted or a request made: file_name must be a bare file name, as it names an entry in the flat cache directory, and without it the first URL’s path must end in one; a keyword the torch function does not take raises TypeError there and then, rather than being mistaken for a corrupt cache entry. A leading ~ in model_dir is expanded.

Parameters:

url (str | list[str]) – a URL string, or a list of URL strings tried left-to-right. A file:// URL is cached by its base name like any other URL, so a cache entry of that name is returned instead of the local file; load a local file with torch.load().
timeout (float | None, optional) – seconds a connection attempt or a single read may stall before the attempt fails; it bounds each wait, not the whole transfer. None uses the KORNIA_DOWNLOAD_TIMEOUT environment variable, read at each call, or 30 s when it is unset. The variable is how to raise the bound for callers that take no timeout, such as pretrained=True model constructors. Default: None
**kwargs (Any) – forwarded to torch.hub.load_state_dict_from_url() (map_location, check_hash, file_name, …), with weights_only set to True unless it is passed as False.

Return type:

dict[str, Any]

Returns:

The loaded state dict.

Raises:

TypeError – if url is not a URL string or a list of them, timeout is not a number, or a keyword is not one the torch function takes.
ValueError – if url is empty or has no scheme (a local path is not a URL); if timeout is not greater than 0 and at most threading.TIMEOUT_MAX, or it is omitted and KORNIA_DOWNLOAD_TIMEOUT holds such a value or no number (the message names which); if file_name is not a bare file name, or it is omitted and the first URL does not end in one.
RuntimeError – if every URL fails. The message carries the failing exception’s type and text, the source it came from and the cache path in play, and the exception itself is chained; without them a rate limit is indistinguishable from a dead link in a CI failure summary, and a caller stuck behind a corrupt entry has no file to delete. Where a load failure set the cache entry aside and the refetch then failed too, the load failure is the one reported and chained, and the refetch failure rides along as context.

Example

>>> sd = load_state_dict_from_url([
...     hf_url("hardnet", "HardNetPP.pth"),  # primary  (HF mirror)
...     "https://github.com/DagnyT/hardnet/raw/master/"
...     "pretrained/pretrained_all_datasets/HardNet%2B%2B.pth",  # fallback
... ])

kornia.core.download_file_from_url(url, *, file_name=None, model_dir=None, progress=True, validate=None, timeout=None)[source]#

Download a file into the torch hub cache and return its path, without loading it.

The sibling of load_state_dict_from_url() for a checkpoint torch cannot unpickle – a .safetensors file, read afterwards with kornia.core.load_safetensors(). It shares that function’s cache, its fallback-URL handling, its retry-with-backoff on transient failures and rate limits (see _MAX_ATTEMPTS and _retry_delay()), its bound on a stalled server (see timeout), and its habit of announcing a transfer on sys.stderr rather than stdout. A file already in the cache is returned as it is, with no request made.

Pass validate to get load_state_dict_from_url()’s quarantine as well. That function hooks its quarantine on the load step, which a download-only function does not have: without one, nothing here can tell a truncated cache entry from an intact one, so a bad file is handed back as a cache hit on every later call and the caller keeps failing until it is deleted by hand.

validate supplies the missing step: it is called with the cache path after each attempt, and raising from it rejects the file. What follows depends on where the file came from.

A rejected cache entry – one the call found rather than fetched – is quarantined as load_state_dict_from_url() quarantines a checkpoint that fails to load: it is moved aside so the remaining sources see an empty path, the source that was handed it is re-fetched once, and if nothing usable turns up the original is put back.

A rejected fresh transfer is deleted outright, whether or not a later source could have used the emptied path. Those bytes arrived during this call and validate refused them, which is a verdict on the file itself, unlike the ambiguous load failures that function has to weigh; keeping them would end the call having added a poisoned entry to a cache that had none. Nothing is re-fetched in that case, so a one-URL call – what download_hf_file() makes – raises with the cache left empty rather than poisoned.

Without validate nothing is quarantined, as before. The path is returned to the caller and named in the failure message so that a file which turns out to be unreadable can be deleted; kornia.core.load_safetensors() names it in every error it raises for the same reason.

Parameters:

url (str | list[str]) – a URL string, or a list of URL strings tried left-to-right. A file:// URL is cached by its base name like any other URL, so a cache entry of that name is returned instead of the local file; read a local file directly.
file_name (str | None, optional) – name to cache the file under. Defaults to the basename of the URL – of the first URL when several are given, so that every source shares one cache slot. Pass it explicitly whenever that basename is not unique to this file: two repositories publishing a model.safetensors each would otherwise share one cache entry, and the second model would silently load the first one’s weights (see _hf_cache_file_name()). Must be a bare filename: the cache is one flat directory, so a value carrying a path separator or naming ./.. would write outside it. Without it, the first URL’s path must end in a file name. Default: None
model_dir (str | None, optional) – directory to cache the file in. Defaults to torch’s <hub dir>/checkpoints, which is the cache CI restores. A leading ~ is expanded. Default: None
progress (bool, optional) – whether to display a progress bar during a transfer. Default: True
validate (Callable[[str], None] | None, optional) – called with the cache path after each attempt, to decide whether what is there is usable. Raising rejects the entry. Keep it cheap – a header parse, not a full read – since it runs on cache hits too. Default: None
timeout (float | None, optional) – seconds a connection attempt or a single read may stall before the attempt fails, which then counts as transient; it bounds each wait, not the whole transfer. None uses the KORNIA_DOWNLOAD_TIMEOUT environment variable, read at each call, or 30 s when it is unset. Default: None

Return type:

str

Returns:

The path of the cached file.

Raises:

TypeError – if url is not a URL string or a list of them, or timeout is not a number.
ValueError – if file_name is not a single path component, or it is omitted and the first URL’s path does not end in a file name; if url is empty or has no scheme (a local path is not a URL); or if timeout is not greater than 0 and at most threading.TIMEOUT_MAX, or it is omitted and KORNIA_DOWNLOAD_TIMEOUT holds such a value or no number (the message names which).
RuntimeError – if every URL fails. The message carries the last failure’s type and text, the source it came from and the cache path in play, and the exception itself is chained.

Example

>>> path = download_file_from_url(
...     hf_url("kimi-vl-a3b-instruct-vision", "model.safetensors"),
...     file_name="kornia--kimi-vl-a3b-instruct-vision--model.safetensors",
... )

kornia.core.download_hf_file(repo, filename, *, model_dir=None, progress=True, validate=None, timeout=None)[source]#

Download one file from a HuggingFace repo and return the path it is cached at.

download_file_from_url() with the two decisions a Hub file needs already made: the resolve/main URL, and a cache name that carries the repository id. The second one is not optional – every repository publishing a single-shard checkpoint calls it model.safetensors, and the cache is one flat directory, so caching under the URL basename would serve one repository’s weights for another’s, silently and on every later call.

Parameters:

repo (str) – repository name under the kornia HF org, or a full owner/name id; see hf_url().
filename (str) – file at the root of that repo.
model_dir (str | None, optional) – directory to cache the file in. Defaults to torch’s <hub dir>/checkpoints, which is the cache CI restores. Default: None
progress (bool, optional) – whether to display a progress bar during a transfer. Default: True
validate (Callable[[str], None] | None, optional) – forwarded to download_file_from_url(), which documents it. Pass kornia.core.check_safetensors() for a checkpoint, so that a truncated cache entry is re-fetched rather than handed back on every later call. Default: None
timeout (float | None, optional) – seconds a connection attempt or a single read may stall before the attempt fails; it bounds each wait, not the whole transfer. None uses the KORNIA_DOWNLOAD_TIMEOUT environment variable, read at each call, or 30 s when it is unset. Forwarded to download_file_from_url(). Default: None

Return type:

str

Returns:

The path of the cached file. Read a .safetensors one with kornia.core.load_safetensors().

Raises:

TypeError – if timeout is not a number.
ValueError – if filename carries a path separator; only files at the repository root are supported, because the cache is one flat directory. Also if timeout is not greater than 0 and at most threading.TIMEOUT_MAX, or it is omitted and KORNIA_DOWNLOAD_TIMEOUT holds such a value or no number (the message names which).
RuntimeError – if the download fails; see download_file_from_url().

Example

>>> path = download_hf_file(
...     "kimi-vl-a3b-instruct-vision", "model.safetensors"
... )

kornia.core.load_safetensors(path, device='cpu')[source]#

Load a .safetensors checkpoint into a state dict.

A pure-torch reader for the format described at huggingface/safetensors: eight bytes of little-endian header length, a JSON header of that length naming each tensor’s dtype, shape and byte range, and one contiguous byte buffer those ranges index into. The optional __metadata__ key is ignored.

The file is memory-mapped rather than read whole, so a multi-gigabyte checkpoint is not held in memory twice. The mapping is private (mmap.ACCESS_COPY) rather than read-only: torch.frombuffer() warns on a buffer it cannot write to, and a private mapping is writable without touching the file – nothing here writes to it, so no page is ever copied. Every tensor is copied out of the mapping before it is returned, so the returned state dict owns its storage and the file is closed by the time this function returns.

Parameters:

path (str | PathLike[str]) – the checkpoint to read.
device (str | device, optional) – the device the returned tensors are placed on. Default: "cpu"

Return type:

dict[str, Tensor]

Returns:

The state dict, in the order the header lists it.

Raises:

ValueError – if the file is not a readable safetensors checkpoint – truncated, a header that is not JSON or nests too deeply to decode, an entry naming a dtype this reader does not accept, a shape torch cannot build, a byte range that does not match the tensor it belongs to, or entries that overlap or leave part of the buffer unaccounted for. Every message names the file.
OSError – if the file cannot be opened or mapped.

Example

>>> state_dict = load_safetensors("model.safetensors")

kornia.core.check_safetensors(path)[source]#

Check that path is a readable safetensors file, without reading tensors.

The header half of load_safetensors(): the length prefix, the JSON header, every entry’s dtype, shape and byte range, and that those ranges cover the buffer exactly once. A file this accepts is one load_safetensors() can build on the CPU, including an empty tensor’s shape, which no byte range bounds; another device can add limits of its own. It reads only the header, so it stays cheap on a multi-gigabyte checkpoint.

Written for kornia.core.download_file_from_url()’s validate argument, where it turns a truncated cache entry into a re-download instead of a permanent failure. Every truncation the download path can produce is visible here: a file cut short after a 2xx leaves the declared header length or the entries’ data_offsets pointing past the end.

Parameters:

path (str | PathLike[str]) – the checkpoint to check.

Raises:

ValueError – if the file is not a readable safetensors checkpoint. The message names path, as load_safetensors() does.
OSError – if the file cannot be opened.

Return type:

None

Exceptions#

The errors raised by kornia’s KORNIA_CHECK* validation helpers in kornia.core.check (and by a few functions directly), such as a tensor with the wrong shape. All of them derive from BaseError, so except BaseError catches any of them. Each subclass also derives from the built-in exception that matches its failure: TypeCheckError is also a TypeError, and ShapeError, ValueCheckError, DeviceError and ImageError are each also a ValueError. Many functions validate their inputs themselves and raise the built-in ValueError or TypeError, so except ValueError or except TypeError catches those and the matching kornia errors alike; except BaseError does not catch those built-ins. BaseError itself derives only from Exception: a failed KORNIA_CHECK(condition) raises it, and neither of those two clauses catches it.

exception kornia.core.exceptions.BaseError[source]#

Bases: Exception

Base class of the exceptions raised by kornia’s KORNIA_CHECK* validation helpers and a few functions.

A failed KORNIA_CHECK raises it directly; it is neither a ValueError nor a TypeError. Each subclass also derives from the one of those two built-ins that matches its failure, so the matching except clause catches the subclass as well.

exception kornia.core.exceptions.ShapeError(message, *, actual_shape=None, expected_shape=None)[source]#