diff --git a/k_diffusion/augmentation.py b/k_diffusion/augmentation.py index 920ef86..917473a 100644 --- a/k_diffusion/augmentation.py +++ b/k_diffusion/augmentation.py @@ -93,7 +93,7 @@ class KarrasAugmentWrapper(nn.Module): def __init__(self, model): super().__init__() self.inner_model = model - + def forward(self, input, sigma, aug_cond=None, mapping_cond=None, **kwargs): if aug_cond is None: aug_cond = input.new_zeros([input.shape[0], 9]) diff --git a/k_diffusion/external.py b/k_diffusion/external.py index 79b51ce..25110de 100644 --- a/k_diffusion/external.py +++ b/k_diffusion/external.py @@ -27,14 +27,14 @@ def t_to_sigma(self, t): return (t * math.pi / 2).tan() def loss(self, input, noise, sigma, **kwargs): - c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_skip, c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) noised_input = input + noise * utils.append_dims(sigma, input.ndim) model_output = self.inner_model(noised_input * c_in, self.sigma_to_t(sigma), **kwargs) target = (input - c_skip * noised_input) / c_out return (model_output - target).pow(2).flatten(1).mean(1) def forward(self, input, sigma, **kwargs): - c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_skip, c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) return self.inner_model(input * c_in, self.sigma_to_t(sigma), **kwargs) * c_out + input * c_skip @@ -102,13 +102,13 @@ def get_eps(self, *args, **kwargs): return self.inner_model(*args, **kwargs) def loss(self, input, noise, sigma, **kwargs): - c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) noised_input = input + noise * utils.append_dims(sigma, input.ndim) eps = self.get_eps(noised_input * c_in, self.sigma_to_t(sigma), **kwargs) return (eps - noise).pow(2).flatten(1).mean(1) def forward(self, input, sigma, **kwargs): - c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) eps = self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs) return input + eps * c_out @@ -156,14 +156,14 @@ def get_v(self, *args, **kwargs): return self.inner_model(*args, **kwargs) def loss(self, input, noise, sigma, **kwargs): - c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_skip, c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) noised_input = input + noise * utils.append_dims(sigma, input.ndim) model_output = self.get_v(noised_input * c_in, self.sigma_to_t(sigma), **kwargs) target = (input - c_skip * noised_input) / c_out return (model_output - target).pow(2).flatten(1).mean(1) def forward(self, input, sigma, **kwargs): - c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_skip, c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) return self.get_v(input * c_in, self.sigma_to_t(sigma), **kwargs) * c_out + input * c_skip diff --git a/k_diffusion/layers.py b/k_diffusion/layers.py index dfe2a0e..34d7a4d 100644 --- a/k_diffusion/layers.py +++ b/k_diffusion/layers.py @@ -73,7 +73,7 @@ def get_scalings(self, sigma): return c_skip, c_out, c_in def loss(self, input, noise, sigma, **kwargs): - c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_skip, c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) c_weight = self.weighting(sigma) noised_input = input + noise * utils.append_dims(sigma, input.ndim) model_output = self.inner_model(noised_input * c_in, sigma, **kwargs) @@ -85,13 +85,13 @@ def loss(self, input, noise, sigma, **kwargs): return (sq_error * f_weight).flatten(1).mean(1) * c_weight def forward(self, input, sigma, **kwargs): - c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_skip, c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) return self.inner_model(input * c_in, sigma, **kwargs) * c_out + input * c_skip class DenoiserWithVariance(Denoiser): def loss(self, input, noise, sigma, **kwargs): - c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_skip, c_out, c_in = (utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)) noised_input = input + noise * utils.append_dims(sigma, input.ndim) model_output, logvar = self.inner_model(noised_input * c_in, sigma, return_variance=True, **kwargs) logvar = utils.append_dims(logvar, model_output.ndim) @@ -234,10 +234,10 @@ def forward(self, input, cond): _kernels = { 'linear': [1 / 8, 3 / 8, 3 / 8, 1 / 8], - 'cubic': + 'cubic': [-0.01171875, -0.03515625, 0.11328125, 0.43359375, 0.43359375, 0.11328125, -0.03515625, -0.01171875], - 'lanczos3': + 'lanczos3': [0.003689131001010537, 0.015056144446134567, -0.03399861603975296, -0.066637322306633, 0.13550527393817902, 0.44638532400131226, 0.44638532400131226, 0.13550527393817902, -0.066637322306633, diff --git a/pyproject.toml b/pyproject.toml index fed528d..eca1e82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,18 @@ [build-system] requires = ["setuptools"] build-backend = "setuptools.build_meta" + +[tool.ruff] +target-version = "py38" + +[tool.ruff.lint] +select = [ + "B", + "E", + "F", + "W", + "UP", +] +ignore = [ + "E501", +]