Skip to content

Commit 3aed05c

Browse files
sayakpaulsayakpaul
committed
Merge branch 'main' into fix/lora-training
2 parents 724ddf9 + 56b3b21 commit 3aed05c

25 files changed

+559
-64
lines changed

‎benchmarks/base_classes.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,25 @@ def run_inference(self, pipe, args):
162162
guidance_scale=1.0,
163163
)
164164

165+
def benchmark(self, args):
166+
flush()
167+
168+
print(f"[INFO] {self.pipe.__class__.__name__}: Running benchmark with: {vars(args)}\n")
169+
170+
time = benchmark_fn(self.run_inference, self.pipe, args) # in seconds.
171+
memory = bytes_to_giga_bytes(torch.cuda.max_memory_allocated()) # in GBs.
172+
benchmark_info = BenchmarkInfo(time=time, memory=memory)
173+
174+
pipeline_class_name = str(self.pipe.__class__.__name__)
175+
flush()
176+
csv_dict = generate_csv_dict(
177+
pipeline_cls=pipeline_class_name, ckpt=self.lora_id, args=args, benchmark_info=benchmark_info
178+
)
179+
filepath = self.get_result_filepath(args)
180+
write_to_csv(filepath, csv_dict)
181+
print(f"Logs written to: {filepath}")
182+
flush()
183+
165184

166185
class ImageToImageBenchmark(TextToImageBenchmark):
167186
pipeline_class = AutoPipelineForImage2Image

‎docs/source/en/api/models/asymmetricautoencoderkl.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,12 @@ make_image_grid([original_image, mask_image, image], rows=1, cols=3)
4949

5050
## AsymmetricAutoencoderKL
5151

52-
[[autodoc]] models.autoencoder_asym_kl.AsymmetricAutoencoderKL
52+
[[autodoc]] models.autoencoders.autoencoder_asym_kl.AsymmetricAutoencoderKL
5353

5454
## AutoencoderKLOutput
5555

56-
[[autodoc]] models.autoencoder_kl.AutoencoderKLOutput
56+
[[autodoc]] models.autoencoders.autoencoder_kl.AutoencoderKLOutput
5757

5858
## DecoderOutput
5959

60-
[[autodoc]] models.vae.DecoderOutput
60+
[[autodoc]] models.autoencoders.vae.DecoderOutput

‎docs/source/en/api/models/autoencoder_tiny.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,4 +54,4 @@ image
5454

5555
## AutoencoderTinyOutput
5656

57-
[[autodoc]] models.autoencoder_tiny.AutoencoderTinyOutput
57+
[[autodoc]] models.autoencoders.autoencoder_tiny.AutoencoderTinyOutput

‎docs/source/en/api/models/autoencoderkl.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,11 +36,11 @@ model = AutoencoderKL.from_single_file(url)
3636

3737
## AutoencoderKLOutput
3838

39-
[[autodoc]] models.autoencoder_kl.AutoencoderKLOutput
39+
[[autodoc]] models.autoencoders.autoencoder_kl.AutoencoderKLOutput
4040

4141
## DecoderOutput
4242

43-
[[autodoc]] models.vae.DecoderOutput
43+
[[autodoc]] models.autoencoders.vae.DecoderOutput
4444

4545
## FlaxAutoencoderKL
4646

‎scripts/convert_consistency_decoder.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@
1212
from tqdm import tqdm
1313

1414
from diffusers import AutoencoderKL, ConsistencyDecoderVAE, DiffusionPipeline, StableDiffusionPipeline, UNet2DModel
15+
from diffusers.models.autoencoders.vae import Encoder
1516
from diffusers.models.embeddings import TimestepEmbedding
1617
from diffusers.models.unet_2d_blocks import ResnetDownsampleBlock2D, ResnetUpsampleBlock2D, UNetMidBlock2D
17-
from diffusers.models.vae import Encoder
1818

1919

2020
args = ArgumentParser()

‎src/diffusers/loaders/lora.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import safetensors
1919
import torch
2020
from huggingface_hub import model_info
21+
from huggingface_hub.constants import HF_HUB_OFFLINE
2122
from huggingface_hub.utils import validate_hf_hub_args
2223
from packaging import version
2324
from torch import nn
@@ -229,7 +230,9 @@ def lora_state_dict(
229230
# determine `weight_name`.
230231
if weight_name is None:
231232
weight_name = cls._best_guess_weight_name(
232-
pretrained_model_name_or_path_or_dict, file_extension=".safetensors"
233+
pretrained_model_name_or_path_or_dict,
234+
file_extension=".safetensors",
235+
local_files_only=local_files_only,
233236
)
234237
model_file = _get_model_file(
235238
pretrained_model_name_or_path_or_dict,
@@ -255,7 +258,7 @@ def lora_state_dict(
255258
if model_file is None:
256259
if weight_name is None:
257260
weight_name = cls._best_guess_weight_name(
258-
pretrained_model_name_or_path_or_dict, file_extension=".bin"
261+
pretrained_model_name_or_path_or_dict, file_extension=".bin", local_files_only=local_files_only
259262
)
260263
model_file = _get_model_file(
261264
pretrained_model_name_or_path_or_dict,
@@ -294,7 +297,12 @@ def lora_state_dict(
294297
return state_dict, network_alphas
295298

296299
@classmethod
297-
def _best_guess_weight_name(cls, pretrained_model_name_or_path_or_dict, file_extension=".safetensors"):
300+
def _best_guess_weight_name(
301+
cls, pretrained_model_name_or_path_or_dict, file_extension=".safetensors", local_files_only=False
302+
):
303+
if local_files_only or HF_HUB_OFFLINE:
304+
raise ValueError("When using the offline mode, you must specify a `weight_name`.")
305+
298306
targeted_files = []
299307

300308
if os.path.isfile(pretrained_model_name_or_path_or_dict):

‎src/diffusers/models/__init__.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,11 @@
2626

2727
if is_torch_available():
2828
_import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"]
29-
_import_structure["autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"]
30-
_import_structure["autoencoder_kl"] = ["AutoencoderKL"]
31-
_import_structure["autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"]
32-
_import_structure["autoencoder_tiny"] = ["AutoencoderTiny"]
33-
_import_structure["consistency_decoder_vae"] = ["ConsistencyDecoderVAE"]
29+
_import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"]
30+
_import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"]
31+
_import_structure["autoencoders.autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"]
32+
_import_structure["autoencoders.autoencoder_tiny"] = ["AutoencoderTiny"]
33+
_import_structure["autoencoders.consistency_decoder_vae"] = ["ConsistencyDecoderVAE"]
3434
_import_structure["controlnet"] = ["ControlNetModel"]
3535
_import_structure["controlnetxs"] = ["ControlNetXSModel"]
3636
_import_structure["dual_transformer_2d"] = ["DualTransformer2DModel"]
@@ -58,11 +58,13 @@
5858
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
5959
if is_torch_available():
6060
from .adapter import MultiAdapter, T2IAdapter
61-
from .autoencoder_asym_kl import AsymmetricAutoencoderKL
62-
from .autoencoder_kl import AutoencoderKL
63-
from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder
64-
from .autoencoder_tiny import AutoencoderTiny
65-
from .consistency_decoder_vae import ConsistencyDecoderVAE
61+
from .autoencoders import (
62+
AsymmetricAutoencoderKL,
63+
AutoencoderKL,
64+
AutoencoderKLTemporalDecoder,
65+
AutoencoderTiny,
66+
ConsistencyDecoderVAE,
67+
)
6668
from .controlnet import ControlNetModel
6769
from .controlnetxs import ControlNetXSModel
6870
from .dual_transformer_2d import DualTransformer2DModel

‎src/diffusers/models/autoencoders/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
from .autoencoder_asym_kl import AsymmetricAutoencoderKL
2+
from .autoencoder_kl import AutoencoderKL
3+
from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder
4+
from .autoencoder_tiny import AutoencoderTiny
5+
from .consistency_decoder_vae import ConsistencyDecoderVAE

src/diffusers/models/autoencoder_asym_kl.py renamed to src/diffusers/models/autoencoders/autoencoder_asym_kl.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,10 @@
1616
import torch
1717
import torch.nn as nn
1818

19-
from ..configuration_utils import ConfigMixin, register_to_config
20-
from ..utils.accelerate_utils import apply_forward_hook
21-
from .modeling_outputs import AutoencoderKLOutput
22-
from .modeling_utils import ModelMixin
19+
from ...configuration_utils import ConfigMixin, register_to_config
20+
from ...utils.accelerate_utils import apply_forward_hook
21+
from ..modeling_outputs import AutoencoderKLOutput
22+
from ..modeling_utils import ModelMixin
2323
from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder, MaskConditionDecoder
2424

2525

src/diffusers/models/autoencoder_kl.py renamed to src/diffusers/models/autoencoders/autoencoder_kl.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,19 @@
1616
import torch
1717
import torch.nn as nn
1818

19-
from ..configuration_utils import ConfigMixin, register_to_config
20-
from ..loaders import FromOriginalVAEMixin
21-
from ..utils.accelerate_utils import apply_forward_hook
22-
from .attention_processor import (
19+
from ...configuration_utils import ConfigMixin, register_to_config
20+
from ...loaders import FromOriginalVAEMixin
21+
from ...utils.accelerate_utils import apply_forward_hook
22+
from ..attention_processor import (
2323
ADDED_KV_ATTENTION_PROCESSORS,
2424
CROSS_ATTENTION_PROCESSORS,
2525
Attention,
2626
AttentionProcessor,
2727
AttnAddedKVProcessor,
2828
AttnProcessor,
2929
)
30-
from .modeling_outputs import AutoencoderKLOutput
31-
from .modeling_utils import ModelMixin
30+
from ..modeling_outputs import AutoencoderKLOutput
31+
from ..modeling_utils import ModelMixin
3232
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder
3333

3434

0 commit comments

Comments
 (0)