From 081d4cf106821793de5ab176c4b08f2b879c0e30 Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Tue, 18 Jun 2024 08:45:57 -0700 Subject: [PATCH] only square images for now --- pyproject.toml | 2 +- titok_pytorch/titok.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c20f6e0..3b23357 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "titok-pytorch" -version = "0.0.1" +version = "0.0.2" description = "TiTok - Pytorch" authors = [ { name = "Phil Wang", email = "lucidrains@gmail.com" } diff --git a/titok_pytorch/titok.py b/titok_pytorch/titok.py index c50110f..6e25cf6 100644 --- a/titok_pytorch/titok.py +++ b/titok_pytorch/titok.py @@ -61,6 +61,8 @@ def __init__( l - latents """ + self.image_size = image_size + assert divisible_by(image_size, patch_size) dim_patch = channels * patch_size ** 2 @@ -139,6 +141,8 @@ def forward( return_codebook_ids = False, return_recons = False ): + assert images.ndim == 4 and images.shape[-2:] == ((self.image_size,) * 2) + batch = images.shape[0] orig_images = images