Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions examples/boft_controlnet/utils/light_controlnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@


from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union
from typing import Optional, Union

import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
Expand All @@ -34,7 +34,7 @@

@dataclass
class ControlNetOutput(BaseOutput):
down_block_res_samples: Tuple[torch.Tensor]
down_block_res_samples: tuple[torch.Tensor]
mid_block_res_sample: torch.Tensor


Expand All @@ -52,7 +52,7 @@ def __init__(
self,
conditioning_embedding_channels: int,
conditioning_channels: int = 3,
block_out_channels: Tuple[int] = (16, 32, 96, 256),
block_out_channels: tuple[int] = (16, 32, 96, 256),
):
super().__init__()

Expand Down Expand Up @@ -92,7 +92,7 @@ def __init__(
in_channels: int = 4,
out_channels: int = 320,
controlnet_conditioning_channel_order: str = "rgb",
conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),
conditioning_embedding_out_channels: Optional[tuple[int]] = (16, 32, 96, 256),
):
super().__init__()

Expand All @@ -104,7 +104,7 @@ def __init__(

@property
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
def attn_processors(self) -> Dict[str, AttentionProcessor]:
def attn_processors(self) -> dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
Expand All @@ -113,7 +113,7 @@ def attn_processors(self) -> Dict[str, AttentionProcessor]:
# set recursively
processors = {}

def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict[str, AttentionProcessor]):
if hasattr(module, "set_processor"):
processors[f"{name}.processor"] = module.processor

Expand All @@ -128,7 +128,7 @@ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors:
return processors

# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
def set_attn_processor(self, processor: Union[AttentionProcessor, dict[str, AttentionProcessor]]):
r"""
Parameters:
`processor (`dict` of `AttentionProcessor` or `AttentionProcessor`):
Expand Down Expand Up @@ -220,7 +220,7 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
# Recursively walk through all the children.
# Any children which exposes the set_attention_slice method
# gets the message
def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: list[int]):
if hasattr(module, "set_attention_slice"):
module.set_attention_slice(slice_size.pop())

Expand All @@ -238,7 +238,7 @@ def _set_gradient_checkpointing(self, module, value=False):
def forward(
self,
controlnet_cond: torch.FloatTensor,
) -> Union[ControlNetOutput, Tuple]:
) -> Union[ControlNetOutput, tuple]:
# check channel order
channel_order = self.config.controlnet_conditioning_channel_order

Expand Down
22 changes: 11 additions & 11 deletions examples/boft_controlnet/utils/pipeline_controlnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# limitations under the License.

from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Union
from typing import Any, Callable, Optional, Union

import numpy as np
import PIL.Image
Expand Down Expand Up @@ -42,8 +42,8 @@ class LightControlNetPipelineOutput(BaseOutput):
(nsfw) content, or `None` if safety checking could not be performed.
"""

images: Union[List[PIL.Image.Image], np.ndarray]
nsfw_content_detected: Optional[List[bool]]
images: Union[list[PIL.Image.Image], np.ndarray]
nsfw_content_detected: Optional[list[bool]]


class LightControlNetPipeline(StableDiffusionControlNetPipeline):
Expand Down Expand Up @@ -164,32 +164,32 @@ def check_inputs(
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
prompt: Union[str, list[str]] = None,
image: Union[
torch.FloatTensor,
PIL.Image.Image,
np.ndarray,
List[torch.FloatTensor],
List[PIL.Image.Image],
List[np.ndarray],
list[torch.FloatTensor],
list[PIL.Image.Image],
list[np.ndarray],
] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
negative_prompt: Optional[Union[str, List[str]]] = None,
negative_prompt: Optional[Union[str, list[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: int = 1,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
cross_attention_kwargs: Optional[dict[str, Any]] = None,
controlnet_conditioning_scale: Union[float, list[float]] = 1.0,
guess_mode: bool = False,
):
r"""
Expand Down
10 changes: 5 additions & 5 deletions examples/boft_controlnet/utils/unet_2d_condition.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# limitations under the License.

from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Union
from typing import Any, Optional, Union

import torch
from diffusers.models import UNet2DConditionModel
Expand Down Expand Up @@ -44,13 +44,13 @@ def forward(
class_labels: Optional[torch.Tensor] = None,
timestep_cond: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
cross_attention_kwargs: Optional[dict[str, Any]] = None,
added_cond_kwargs: Optional[dict[str, torch.Tensor]] = None,
down_block_additional_residuals: Optional[tuple[torch.Tensor]] = None,
mid_block_additional_residual: Optional[torch.Tensor] = None,
encoder_attention_mask: Optional[torch.Tensor] = None,
return_dict: bool = True,
) -> Union[UNet2DConditionOutput, Tuple]:
) -> Union[UNet2DConditionOutput, tuple]:
r"""
Args:
sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
Expand Down
13 changes: 7 additions & 6 deletions examples/corda_finetuning/corda_finetuning.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,9 @@

import copy
import os
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Sequence
from typing import Optional

import torch
import transformers
Expand Down Expand Up @@ -65,7 +66,7 @@ class TrainingArguments(transformers.TrainingArguments):
model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
data_path: str = field(default=None, metadata={"help": "Path to the training data."})
dataset_split: str = field(default="train[:100000]", metadata={"help": "(`['train', 'test', 'eval']`):"})
dataset_field: List[str] = field(default=None, metadata={"help": "Fields of dataset input and output."})
dataset_field: list[str] = field(default=None, metadata={"help": "Fields of dataset input and output."})
dataloader_num_proc: int = field(default=16, metadata={"help": "Number of processes to load dataset"})
dataloader_batch_size: int = field(
default=3000,
Expand Down Expand Up @@ -95,7 +96,7 @@ def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: st


def smart_tokenizer_and_embedding_resize(
special_tokens_dict: Dict,
special_tokens_dict: dict,
tokenizer: transformers.PreTrainedTokenizer,
model: transformers.PreTrainedModel,
):
Expand All @@ -117,7 +118,7 @@ def smart_tokenizer_and_embedding_resize(
output_embeddings[-num_new_tokens:] = output_embeddings_avg


def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> dict:
"""Tokenize a list of strings."""
tokenized_list = [
tokenizer(
Expand Down Expand Up @@ -145,7 +146,7 @@ def preprocess(
sources: Sequence[str],
targets: Sequence[str],
tokenizer: transformers.PreTrainedTokenizer,
) -> Dict:
) -> dict:
"""Preprocess the data by tokenizing."""
examples = [s + t for s, t in zip(sources, targets)]
examples_tokenized, sources_tokenized = (_tokenize_fn(strings, tokenizer) for strings in (examples, sources))
Expand All @@ -165,7 +166,7 @@ class DataCollatorForSupervisedDataset:

tokenizer: transformers.PreTrainedTokenizer

def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
def __call__(self, instances: Sequence[dict]) -> dict[str, torch.Tensor]:
input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
input_ids = [torch.tensor(x) for x in input_ids]
input_ids = torch.nn.utils.rnn.pad_sequence(
Expand Down
4 changes: 2 additions & 2 deletions examples/int8_training/peft_adalora_whisper_large_training.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from datetime import datetime
from pathlib import Path
from random import randint
from typing import Any, Dict, List, Union
from typing import Any, Union

# datasets imports
import datasets
Expand Down Expand Up @@ -337,7 +337,7 @@ def load_model_hook(models, input_dir):
class DataCollatorSpeechSeq2SeqWithPadding:
processor: Any

def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
def __call__(self, features: list[dict[str, Union[list[int], torch.Tensor]]]) -> dict[str, torch.Tensor]:
# split inputs and labels since they have to be of different lengths and need different padding methods
# first treat the audio inputs by simply returning torch tensors
input_features = [{"input_features": feature["input_features"]} for feature in features]
Expand Down
10 changes: 5 additions & 5 deletions examples/lora_dreambooth/convert_kohya_ss_sd_lora_to_peft.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import os
from collections import Counter
from dataclasses import dataclass
from typing import Dict, Optional
from typing import Optional

import safetensors
import torch
Expand Down Expand Up @@ -30,13 +30,13 @@ class LoRAInfo:
lora_A: Optional[torch.Tensor] = None
lora_B: Optional[torch.Tensor] = None

def peft_state_dict(self) -> Dict[str, torch.Tensor]:
def peft_state_dict(self) -> dict[str, torch.Tensor]:
if self.lora_A is None or self.lora_B is None:
raise ValueError("At least one of lora_A or lora_B is None, they must both be provided")
return {f"{peft_key}.lora_A.weight": self.lora_A, f"{peft_key}.lora_B.weight": self.lora_A}


def construct_peft_loraconfig(info: Dict[str, LoRAInfo]) -> LoraConfig:
def construct_peft_loraconfig(info: dict[str, LoRAInfo]) -> LoraConfig:
"""Constructs LoraConfig from data extracted from kohya checkpoint

Args:
Expand Down Expand Up @@ -75,7 +75,7 @@ def construct_peft_loraconfig(info: Dict[str, LoRAInfo]) -> LoraConfig:
return config


def combine_peft_state_dict(info: Dict[str, LoRAInfo]) -> Dict[str, torch.Tensor]:
def combine_peft_state_dict(info: dict[str, LoRAInfo]) -> dict[str, torch.Tensor]:
result = {}
for key_name, key_info in info.items():
result[f"base_model.model.{key_name}.lora_A.weight"] = key_info.lora_A
Expand Down Expand Up @@ -115,7 +115,7 @@ def combine_peft_state_dict(info: Dict[str, LoRAInfo]) -> Dict[str, torch.Tensor
)

# Store conversion info (model_type -> peft_key -> LoRAInfo)
lora_info: Dict[str, Dict[str, LoRAInfo]] = {
lora_info: dict[str, dict[str, LoRAInfo]] = {
"text_encoder": {},
"unet": {},
}
Expand Down
3 changes: 1 addition & 2 deletions examples/lora_dreambooth/convert_peft_sd_lora_to_kohya_ss.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import argparse
import os
from typing import Dict

import torch
from diffusers import UNet2DConditionModel
Expand All @@ -19,7 +18,7 @@

def get_module_kohya_state_dict(
module: PeftModel, prefix: str, dtype: torch.dtype, adapter_name: str = LORA_ADAPTER_NAME
) -> Dict[str, torch.Tensor]:
) -> dict[str, torch.Tensor]:
kohya_ss_state_dict = {}
for peft_key, weight in get_peft_model_state_dict(module, adapter_name=adapter_name).items():
kohya_key = peft_key.replace("base_model.model", prefix)
Expand Down
4 changes: 2 additions & 2 deletions examples/olora_finetuning/olora_finetuning.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@


import os
from typing import List, Optional
from typing import Optional

import torch
import transformers
Expand Down Expand Up @@ -43,7 +43,7 @@ def train(
lora_r: int = 32,
lora_alpha: int = 16,
lora_dropout: float = 0.05,
lora_target_modules: List[str] = None,
lora_target_modules: list[str] = None,
torch_dtype: str = "float16",
init_lora_weights="olora",
seed: Optional[int] = None,
Expand Down
Loading