From 9d2fbf6dbd1aee10a8b5a4f79b3f26f33e7ea70e Mon Sep 17 00:00:00 2001 From: limiteinductive Date: Thu, 31 Aug 2023 12:40:04 +0200 Subject: [PATCH] Fix tuple annotation for pyright 1.1.325 --- src/refiners/fluxion/layers/converter.py | 2 +- src/refiners/training_utils/latent_diffusion.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/refiners/fluxion/layers/converter.py b/src/refiners/fluxion/layers/converter.py index 653b3a8..653bb74 100644 --- a/src/refiners/fluxion/layers/converter.py +++ b/src/refiners/fluxion/layers/converter.py @@ -24,7 +24,7 @@ class Converter(ContextModule): self.set_device = set_device self.set_dtype = set_dtype - def forward(self, *inputs: Tensor) -> tuple[Tensor]: + def forward(self, *inputs: Tensor) -> tuple[Tensor, ...]: parent = self.ensure_parent converted_tensors: list[Tensor] = [] diff --git a/src/refiners/training_utils/latent_diffusion.py b/src/refiners/training_utils/latent_diffusion.py index bdd75be..4ce4d4a 100644 --- a/src/refiners/training_utils/latent_diffusion.py +++ b/src/refiners/training_utils/latent_diffusion.py @@ -160,7 +160,7 @@ class LatentDiffusionTrainer(Trainer[ConfigType, TextEmbeddingLatentsBatch]): self.current_step = random_step return self.ddpm_scheduler.timesteps[random_step].unsqueeze(dim=0) - def sample_noise(self, size: tuple[int, int, int, int], dtype: DType | None = None) -> Tensor: + def sample_noise(self, size: tuple[int, ...], dtype: DType | None = None) -> Tensor: return sample_noise( size=size, offset_noise=self.config.latent_diffusion.offset_noise, device=self.device, dtype=dtype ) @@ -207,7 +207,7 @@ class LatentDiffusionTrainer(Trainer[ConfigType, TextEmbeddingLatentsBatch]): def sample_noise( - size: tuple[int, int, int, int], + size: tuple[int, ...], offset_noise: float = 0.1, device: Device | str = "cpu", dtype: DType | None = None,