ruff 3 formatting (Rye 0.28)

This commit is contained in:
Pierre Chapuis 2024-03-08 10:35:42 +01:00
parent a0be5458b9
commit be2368cf20
13 changed files with 46 additions and 90 deletions

View file

@ -4,6 +4,7 @@ Download and convert weights for testing
To see what weights will be downloaded and converted, run: To see what weights will be downloaded and converted, run:
DRY_RUN=1 python scripts/prepare_test_weights.py DRY_RUN=1 python scripts/prepare_test_weights.py
""" """
import hashlib import hashlib
import os import os
import subprocess import subprocess

View file

@ -129,8 +129,7 @@ class Lora(Generic[T], fl.Chain, ABC):
return loras return loras
@abstractmethod @abstractmethod
def is_compatible(self, layer: fl.WeightedModule, /) -> bool: def is_compatible(self, layer: fl.WeightedModule, /) -> bool: ...
...
def auto_attach( def auto_attach(
self, self,

View file

@ -256,16 +256,13 @@ class Chain(ContextModule):
self._modules = generate_unique_names(tuple(modules)) # type: ignore self._modules = generate_unique_names(tuple(modules)) # type: ignore
@overload @overload
def __getitem__(self, key: int) -> Module: def __getitem__(self, key: int) -> Module: ...
...
@overload @overload
def __getitem__(self, key: str) -> Module: def __getitem__(self, key: str) -> Module: ...
...
@overload @overload
def __getitem__(self, key: slice) -> "Chain": def __getitem__(self, key: slice) -> "Chain": ...
...
def __getitem__(self, key: int | str | slice) -> Module: def __getitem__(self, key: int | str | slice) -> Module:
if isinstance(key, slice): if isinstance(key, slice):
@ -324,16 +321,14 @@ class Chain(ContextModule):
self, self,
predicate: Callable[[Module, "Chain"], bool] | None = None, predicate: Callable[[Module, "Chain"], bool] | None = None,
recurse: bool = False, recurse: bool = False,
) -> Iterator[tuple[Module, "Chain"]]: ) -> Iterator[tuple[Module, "Chain"]]: ...
...
@overload @overload
def walk( def walk(
self, self,
predicate: type[T], predicate: type[T],
recurse: bool = False, recurse: bool = False,
) -> Iterator[tuple[T, "Chain"]]: ) -> Iterator[tuple[T, "Chain"]]: ...
...
def walk( def walk(
self, self,

View file

@ -440,18 +440,15 @@ class IPAdapter(Generic[T], fl.Chain, Adapter[T]):
self.set_context("ip_adapter", {"clip_image_embedding": image_embedding}) self.set_context("ip_adapter", {"clip_image_embedding": image_embedding})
@overload @overload
def compute_clip_image_embedding(self, image_prompt: Tensor, weights: list[float] | None = None) -> Tensor: def compute_clip_image_embedding(self, image_prompt: Tensor, weights: list[float] | None = None) -> Tensor: ...
...
@overload @overload
def compute_clip_image_embedding(self, image_prompt: Image.Image) -> Tensor: def compute_clip_image_embedding(self, image_prompt: Image.Image) -> Tensor: ...
...
@overload @overload
def compute_clip_image_embedding( def compute_clip_image_embedding(
self, image_prompt: list[Image.Image], weights: list[float] | None = None self, image_prompt: list[Image.Image], weights: list[float] | None = None
) -> Tensor: ) -> Tensor: ...
...
def compute_clip_image_embedding( def compute_clip_image_embedding(
self, self,

View file

@ -66,22 +66,18 @@ class LatentDiffusionModel(fl.Module, ABC):
return self.solver.inference_steps return self.solver.inference_steps
@abstractmethod @abstractmethod
def set_unet_context(self, *, timestep: Tensor, clip_text_embedding: Tensor, **_: Tensor) -> None: def set_unet_context(self, *, timestep: Tensor, clip_text_embedding: Tensor, **_: Tensor) -> None: ...
...
@abstractmethod @abstractmethod
def set_self_attention_guidance(self, enable: bool, scale: float = 1.0) -> None: def set_self_attention_guidance(self, enable: bool, scale: float = 1.0) -> None: ...
...
@abstractmethod @abstractmethod
def has_self_attention_guidance(self) -> bool: def has_self_attention_guidance(self) -> bool: ...
...
@abstractmethod @abstractmethod
def compute_self_attention_guidance( def compute_self_attention_guidance(
self, x: Tensor, noise: Tensor, step: int, *, clip_text_embedding: Tensor, **kwargs: Tensor self, x: Tensor, noise: Tensor, step: int, *, clip_text_embedding: Tensor, **kwargs: Tensor
) -> Tensor: ) -> Tensor: ...
...
def forward( def forward(
self, x: Tensor, step: int, *, clip_text_embedding: Tensor, condition_scale: float = 7.5, **kwargs: Tensor self, x: Tensor, step: int, *, clip_text_embedding: Tensor, condition_scale: float = 7.5, **kwargs: Tensor

View file

@ -68,8 +68,7 @@ class MultiDiffusion(Generic[T, D], ABC):
return torch.where(condition=num_updates > 0, input=cumulative_values / num_updates, other=x) return torch.where(condition=num_updates > 0, input=cumulative_values / num_updates, other=x)
@abstractmethod @abstractmethod
def diffuse_target(self, x: Tensor, step: int, target: D) -> Tensor: def diffuse_target(self, x: Tensor, step: int, target: D) -> Tensor: ...
...
@property @property
def steps(self) -> list[int]: def steps(self) -> list[int]:

View file

@ -20,56 +20,38 @@ class CallbackConfig(BaseModel):
class Callback(Generic[T]): class Callback(Generic[T]):
def on_init_begin(self, trainer: T) -> None: def on_init_begin(self, trainer: T) -> None: ...
...
def on_init_end(self, trainer: T) -> None: def on_init_end(self, trainer: T) -> None: ...
...
def on_train_begin(self, trainer: T) -> None: def on_train_begin(self, trainer: T) -> None: ...
...
def on_train_end(self, trainer: T) -> None: def on_train_end(self, trainer: T) -> None: ...
...
def on_epoch_begin(self, trainer: T) -> None: def on_epoch_begin(self, trainer: T) -> None: ...
...
def on_epoch_end(self, trainer: T) -> None: def on_epoch_end(self, trainer: T) -> None: ...
...
def on_batch_begin(self, trainer: T) -> None: def on_batch_begin(self, trainer: T) -> None: ...
...
def on_batch_end(self, trainer: T) -> None: def on_batch_end(self, trainer: T) -> None: ...
...
def on_backward_begin(self, trainer: T) -> None: def on_backward_begin(self, trainer: T) -> None: ...
...
def on_backward_end(self, trainer: T) -> None: def on_backward_end(self, trainer: T) -> None: ...
...
def on_optimizer_step_begin(self, trainer: T) -> None: def on_optimizer_step_begin(self, trainer: T) -> None: ...
...
def on_optimizer_step_end(self, trainer: T) -> None: def on_optimizer_step_end(self, trainer: T) -> None: ...
...
def on_compute_loss_begin(self, trainer: T) -> None: def on_compute_loss_begin(self, trainer: T) -> None: ...
...
def on_compute_loss_end(self, trainer: T) -> None: def on_compute_loss_end(self, trainer: T) -> None: ...
...
def on_evaluate_begin(self, trainer: T) -> None: def on_evaluate_begin(self, trainer: T) -> None: ...
...
def on_evaluate_end(self, trainer: T) -> None: def on_evaluate_end(self, trainer: T) -> None: ...
...
def on_lr_scheduler_step_begin(self, trainer: T) -> None: def on_lr_scheduler_step_begin(self, trainer: T) -> None: ...
...
def on_lr_scheduler_step_end(self, trainer: T) -> None: def on_lr_scheduler_step_end(self, trainer: T) -> None: ...
...

View file

@ -10,11 +10,9 @@ T = TypeVar("T", covariant=True)
class HuggingfaceDataset(Generic[T], Protocol): class HuggingfaceDataset(Generic[T], Protocol):
def __getitem__(self, index: int) -> T: def __getitem__(self, index: int) -> T: ...
...
def __len__(self) -> int: def __len__(self) -> int: ...
...
def load_hf_dataset( def load_hf_dataset(

View file

@ -335,8 +335,7 @@ class Trainer(Generic[ConfigType, Batch], ABC):
) )
@abstractmethod @abstractmethod
def compute_loss(self, batch: Batch) -> Tensor: def compute_loss(self, batch: Batch) -> Tensor: ...
...
def compute_evaluation(self) -> None: def compute_evaluation(self) -> None:
pass pass

View file

@ -10,13 +10,11 @@ from refiners.foundationals.latent_diffusion.image_prompt import ImageCrossAtten
@overload @overload
def new_adapter(target: SD1UNet) -> SD1IPAdapter: def new_adapter(target: SD1UNet) -> SD1IPAdapter: ...
...
@overload @overload
def new_adapter(target: SDXLUNet) -> SDXLIPAdapter: def new_adapter(target: SDXLUNet) -> SDXLIPAdapter: ...
...
def new_adapter(target: SD1UNet | SDXLUNet) -> SD1IPAdapter | SDXLIPAdapter: def new_adapter(target: SD1UNet | SDXLUNet) -> SD1IPAdapter | SDXLIPAdapter:

View file

@ -9,13 +9,11 @@ from refiners.foundationals.latent_diffusion.t2i_adapter import T2IFeatures
@overload @overload
def new_adapter(target: SD1UNet, name: str) -> SD1T2IAdapter: def new_adapter(target: SD1UNet, name: str) -> SD1T2IAdapter: ...
...
@overload @overload
def new_adapter(target: SDXLUNet, name: str) -> SDXLT2IAdapter: def new_adapter(target: SDXLUNet, name: str) -> SDXLT2IAdapter: ...
...
def new_adapter(target: SD1UNet | SDXLUNet, name: str) -> SD1T2IAdapter | SDXLT2IAdapter: def new_adapter(target: SD1UNet | SDXLUNet, name: str) -> SD1T2IAdapter | SDXLT2IAdapter:

View file

@ -19,8 +19,7 @@ class DiffusersSDXL(Protocol):
tokenizer_2: fl.Module tokenizer_2: fl.Module
vae: fl.Module vae: fl.Module
def __call__(self, prompt: str, *args: Any, **kwargs: Any) -> Any: def __call__(self, prompt: str, *args: Any, **kwargs: Any) -> Any: ...
...
def encode_prompt( def encode_prompt(
self, self,
@ -28,8 +27,7 @@ class DiffusersSDXL(Protocol):
prompt_2: str | None = None, prompt_2: str | None = None,
negative_prompt: str | None = None, negative_prompt: str | None = None,
negative_prompt_2: str | None = None, negative_prompt_2: str | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ...
...
@pytest.fixture(scope="module") @pytest.fixture(scope="module")

View file

@ -32,19 +32,16 @@ class FacebookSAM(nn.Module):
prompt_encoder: nn.Module prompt_encoder: nn.Module
mask_decoder: nn.Module mask_decoder: nn.Module
def __call__(self, batched_input: list[SAMInput], multimask_output: bool) -> list[SAMOutput]: def __call__(self, batched_input: list[SAMInput], multimask_output: bool) -> list[SAMOutput]: ...
...
@property @property
def device(self) -> Any: def device(self) -> Any: ...
...
class FacebookSAMPredictor: class FacebookSAMPredictor:
model: FacebookSAM model: FacebookSAM
def set_image(self, image: NDArrayUInt8, image_format: str = "RGB") -> None: def set_image(self, image: NDArrayUInt8, image_format: str = "RGB") -> None: ...
...
def predict( def predict(
self, self,
@ -54,8 +51,7 @@ class FacebookSAMPredictor:
mask_input: NDArray | None = None, mask_input: NDArray | None = None,
multimask_output: bool = True, multimask_output: bool = True,
return_logits: bool = False, return_logits: bool = False,
) -> tuple[NDArray, NDArray, NDArray]: ) -> tuple[NDArray, NDArray, NDArray]: ...
...
@dataclass @dataclass