diff --git a/src/diffusers/pipeline_utils.py b/src/diffusers/pipeline_utils.py index d73b8d8fb3..b0ff25c339 100644 --- a/src/diffusers/pipeline_utils.py +++ b/src/diffusers/pipeline_utils.py @@ -57,17 +57,19 @@ class DiffusionPipeline(ConfigMixin): from diffusers import pipelines for name, module in kwargs.items(): - # check if the module is a pipeline module - is_pipeline_module = hasattr(pipelines, module.__module__.split(".")[-1]) - # retrive library library = module.__module__.split(".")[0] + # check if the module is a pipeline module + pipeline_file = module.__module__.split(".")[-1] + pipeline_dir = module.__module__.split(".")[-2] + is_pipeline_module = pipeline_file == "pipeline_" + pipeline_dir and hasattr(pipelines, pipeline_dir) + # if library is not in LOADABLE_CLASSES, then it is a custom module. # Or if it's a pipeline module, then the module is inside the pipeline - # so we set the library to module name. + # folder so we set the library to module name. if library not in LOADABLE_CLASSES or is_pipeline_module: - library = module.__module__.split(".")[-1] + library = pipeline_dir # retrive class_name class_name = module.__class__.__name__ diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index b956a2b8f8..b8cab072cd 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -1,20 +1,17 @@ from ..utils import is_inflect_available, is_transformers_available, is_unidecode_available -from .pipeline_bddm import BDDMPipeline -from .pipeline_ddim import DDIMPipeline -from .pipeline_ddpm import DDPMPipeline -from .pipeline_latent_diffusion_uncond import LatentDiffusionUncondPipeline -from .pipeline_pndm import PNDMPipeline -from .pipeline_score_sde_ve import ScoreSdeVePipeline -from .pipeline_score_sde_vp import ScoreSdeVpPipeline - - -# from .pipeline_score_sde import ScoreSdeVePipeline +from .bddm import BDDMPipeline +from .ddim import DDIMPipeline +from .ddpm import DDPMPipeline +from .latent_diffusion_uncond import LatentDiffusionUncondPipeline +from .pndm import PNDMPipeline +from .score_sde_ve import ScoreSdeVePipeline +from .score_sde_vp import ScoreSdeVpPipeline if is_transformers_available(): - from .pipeline_glide import GlidePipeline - from .pipeline_latent_diffusion import LatentDiffusionPipeline + from .glide import GlidePipeline + from .latent_diffusion import LatentDiffusionPipeline if is_transformers_available() and is_unidecode_available() and is_inflect_available(): - from .pipeline_grad_tts import GradTTSPipeline + from .grad_tts import GradTTSPipeline diff --git a/src/diffusers/pipelines/bddm/__init__.py b/src/diffusers/pipelines/bddm/__init__.py new file mode 100644 index 0000000000..6505e93100 --- /dev/null +++ b/src/diffusers/pipelines/bddm/__init__.py @@ -0,0 +1 @@ +from .pipeline_bddm import BDDMPipeline, DiffWave diff --git a/src/diffusers/pipelines/pipeline_bddm.py b/src/diffusers/pipelines/bddm/pipeline_bddm.py similarity index 98% rename from src/diffusers/pipelines/pipeline_bddm.py rename to src/diffusers/pipelines/bddm/pipeline_bddm.py index 09120fdab0..45b26c3127 100644 --- a/src/diffusers/pipelines/pipeline_bddm.py +++ b/src/diffusers/pipelines/bddm/pipeline_bddm.py @@ -21,9 +21,9 @@ import torch.nn.functional as F import tqdm -from ..configuration_utils import ConfigMixin -from ..modeling_utils import ModelMixin -from ..pipeline_utils import DiffusionPipeline +from ...configuration_utils import ConfigMixin +from ...modeling_utils import ModelMixin +from ...pipeline_utils import DiffusionPipeline def calc_diffusion_step_embedding(diffusion_steps, diffusion_step_embed_dim_in): diff --git a/src/diffusers/pipelines/ddim/__init__.py b/src/diffusers/pipelines/ddim/__init__.py new file mode 100644 index 0000000000..85e8118e75 --- /dev/null +++ b/src/diffusers/pipelines/ddim/__init__.py @@ -0,0 +1 @@ +from .pipeline_ddim import DDIMPipeline diff --git a/src/diffusers/pipelines/pipeline_ddim.py b/src/diffusers/pipelines/ddim/pipeline_ddim.py similarity index 98% rename from src/diffusers/pipelines/pipeline_ddim.py rename to src/diffusers/pipelines/ddim/pipeline_ddim.py index 8da24dbf8f..80bbb5a622 100644 --- a/src/diffusers/pipelines/pipeline_ddim.py +++ b/src/diffusers/pipelines/ddim/pipeline_ddim.py @@ -18,7 +18,7 @@ import torch import tqdm -from ..pipeline_utils import DiffusionPipeline +from ...pipeline_utils import DiffusionPipeline class DDIMPipeline(DiffusionPipeline): diff --git a/src/diffusers/pipelines/ddpm/__init__.py b/src/diffusers/pipelines/ddpm/__init__.py new file mode 100644 index 0000000000..bb228ee012 --- /dev/null +++ b/src/diffusers/pipelines/ddpm/__init__.py @@ -0,0 +1 @@ +from .pipeline_ddpm import DDPMPipeline diff --git a/src/diffusers/pipelines/pipeline_ddpm.py b/src/diffusers/pipelines/ddpm/pipeline_ddpm.py similarity index 97% rename from src/diffusers/pipelines/pipeline_ddpm.py rename to src/diffusers/pipelines/ddpm/pipeline_ddpm.py index 9cf83bfb75..dd00daa9aa 100644 --- a/src/diffusers/pipelines/pipeline_ddpm.py +++ b/src/diffusers/pipelines/ddpm/pipeline_ddpm.py @@ -18,7 +18,7 @@ import torch import tqdm -from ..pipeline_utils import DiffusionPipeline +from ...pipeline_utils import DiffusionPipeline class DDPMPipeline(DiffusionPipeline): diff --git a/src/diffusers/pipelines/glide/__init__.py b/src/diffusers/pipelines/glide/__init__.py new file mode 100644 index 0000000000..d4bac0ebf6 --- /dev/null +++ b/src/diffusers/pipelines/glide/__init__.py @@ -0,0 +1,5 @@ +from ...utils import is_transformers_available + + +if is_transformers_available(): + from .pipeline_glide import CLIPTextModel, GlidePipeline diff --git a/src/diffusers/pipelines/pipeline_glide.py b/src/diffusers/pipelines/glide/pipeline_glide.py similarity index 99% rename from src/diffusers/pipelines/pipeline_glide.py rename to src/diffusers/pipelines/glide/pipeline_glide.py index 9a67790b35..a9b6746f5c 100644 --- a/src/diffusers/pipelines/pipeline_glide.py +++ b/src/diffusers/pipelines/glide/pipeline_glide.py @@ -18,7 +18,6 @@ import math from dataclasses import dataclass from typing import Any, Optional, Tuple, Union -import numpy as np import torch import torch.utils.checkpoint from torch import nn @@ -30,10 +29,10 @@ from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPo from transformers.modeling_utils import PreTrainedModel from transformers.utils import ModelOutput, add_start_docstrings_to_model_forward, replace_return_docstrings -from ..models import GlideSuperResUNetModel, GlideTextToImageUNetModel -from ..pipeline_utils import DiffusionPipeline -from ..schedulers import DDIMScheduler, DDPMScheduler -from ..utils import logging +from ...models import GlideSuperResUNetModel, GlideTextToImageUNetModel +from ...pipeline_utils import DiffusionPipeline +from ...schedulers import DDIMScheduler, DDPMScheduler +from ...utils import logging ##################### @@ -594,7 +593,7 @@ class CLIPTextTransformer(nn.Module): bsz, seq_len = input_shape # CLIP's text model uses causal mask, prepare it here. # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 - causal_attention_mask = self._build_causal_attention_mask(bsz, seq_len).to(hidden_states.device) + # causal_attention_mask = self._build_causal_attention_mask(bsz, seq_len).to(hidden_states.device) # expand attention_mask if attention_mask is not None: diff --git a/src/diffusers/pipelines/grad_tts/__init__.py b/src/diffusers/pipelines/grad_tts/__init__.py new file mode 100644 index 0000000000..decec2622b --- /dev/null +++ b/src/diffusers/pipelines/grad_tts/__init__.py @@ -0,0 +1,6 @@ +from ...utils import is_inflect_available, is_transformers_available, is_unidecode_available + + +if is_transformers_available() and is_unidecode_available() and is_inflect_available(): + from .grad_tts_utils import GradTTSTokenizer + from .pipeline_grad_tts import GradTTSPipeline, TextEncoder diff --git a/src/diffusers/pipelines/grad_tts_utils.py b/src/diffusers/pipelines/grad_tts/grad_tts_utils.py similarity index 100% rename from src/diffusers/pipelines/grad_tts_utils.py rename to src/diffusers/pipelines/grad_tts/grad_tts_utils.py diff --git a/src/diffusers/pipelines/pipeline_grad_tts.py b/src/diffusers/pipelines/grad_tts/pipeline_grad_tts.py similarity index 99% rename from src/diffusers/pipelines/pipeline_grad_tts.py rename to src/diffusers/pipelines/grad_tts/pipeline_grad_tts.py index 93770fe21e..adf027e10a 100644 --- a/src/diffusers/pipelines/pipeline_grad_tts.py +++ b/src/diffusers/pipelines/grad_tts/pipeline_grad_tts.py @@ -6,10 +6,10 @@ import torch from torch import nn import tqdm -from diffusers import DiffusionPipeline -from diffusers.configuration_utils import ConfigMixin -from diffusers.modeling_utils import ModelMixin +from ...configuration_utils import ConfigMixin +from ...modeling_utils import ModelMixin +from ...pipeline_utils import DiffusionPipeline from .grad_tts_utils import GradTTSTokenizer # flake8: noqa diff --git a/src/diffusers/pipelines/latent_diffusion/__init__.py b/src/diffusers/pipelines/latent_diffusion/__init__.py new file mode 100644 index 0000000000..614c3db4bd --- /dev/null +++ b/src/diffusers/pipelines/latent_diffusion/__init__.py @@ -0,0 +1,5 @@ +from ...utils import is_transformers_available + + +if is_transformers_available(): + from .pipeline_latent_diffusion import AutoencoderKL, LatentDiffusionPipeline, LDMBertModel diff --git a/src/diffusers/pipelines/pipeline_latent_diffusion.py b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py similarity index 99% rename from src/diffusers/pipelines/pipeline_latent_diffusion.py rename to src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py index fea7a287ed..f7dc8ea343 100644 --- a/src/diffusers/pipelines/pipeline_latent_diffusion.py +++ b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py @@ -7,20 +7,15 @@ import torch.nn as nn import torch.utils.checkpoint import tqdm +from transformers.activations import ACT2FN +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_outputs import BaseModelOutput +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging - -try: - from transformers.activations import ACT2FN - from transformers.configuration_utils import PretrainedConfig - from transformers.modeling_outputs import BaseModelOutput - from transformers.modeling_utils import PreTrainedModel - from transformers.utils import logging -except ImportError: - raise ImportError("Please install the transformers.") - -from ..configuration_utils import ConfigMixin -from ..modeling_utils import ModelMixin -from ..pipeline_utils import DiffusionPipeline +from ...configuration_utils import ConfigMixin +from ...modeling_utils import ModelMixin +from ...pipeline_utils import DiffusionPipeline ################################################################################ diff --git a/src/diffusers/pipelines/latent_diffusion_uncond/__init__.py b/src/diffusers/pipelines/latent_diffusion_uncond/__init__.py new file mode 100644 index 0000000000..6fe1f6d655 --- /dev/null +++ b/src/diffusers/pipelines/latent_diffusion_uncond/__init__.py @@ -0,0 +1 @@ +from .pipeline_latent_diffusion_uncond import LatentDiffusionUncondPipeline diff --git a/src/diffusers/pipelines/pipeline_latent_diffusion_uncond.py b/src/diffusers/pipelines/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py similarity index 99% rename from src/diffusers/pipelines/pipeline_latent_diffusion_uncond.py rename to src/diffusers/pipelines/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py index 34348e9111..269dfb39cd 100644 --- a/src/diffusers/pipelines/pipeline_latent_diffusion_uncond.py +++ b/src/diffusers/pipelines/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py @@ -6,9 +6,9 @@ import torch.nn as nn import tqdm -from ..configuration_utils import ConfigMixin -from ..modeling_utils import ModelMixin -from ..pipeline_utils import DiffusionPipeline +from ...configuration_utils import ConfigMixin +from ...modeling_utils import ModelMixin +from ...pipeline_utils import DiffusionPipeline def get_timestep_embedding(timesteps, embedding_dim): diff --git a/src/diffusers/pipelines/pndm/__init__.py b/src/diffusers/pipelines/pndm/__init__.py new file mode 100644 index 0000000000..488eb4f5f2 --- /dev/null +++ b/src/diffusers/pipelines/pndm/__init__.py @@ -0,0 +1 @@ +from .pipeline_pndm import PNDMPipeline diff --git a/src/diffusers/pipelines/pipeline_pndm.py b/src/diffusers/pipelines/pndm/pipeline_pndm.py similarity index 97% rename from src/diffusers/pipelines/pipeline_pndm.py rename to src/diffusers/pipelines/pndm/pipeline_pndm.py index 5fd8a98483..7665b2b593 100644 --- a/src/diffusers/pipelines/pipeline_pndm.py +++ b/src/diffusers/pipelines/pndm/pipeline_pndm.py @@ -18,7 +18,7 @@ import torch import tqdm -from ..pipeline_utils import DiffusionPipeline +from ...pipeline_utils import DiffusionPipeline class PNDMPipeline(DiffusionPipeline): diff --git a/src/diffusers/pipelines/score_sde_ve/__init__.py b/src/diffusers/pipelines/score_sde_ve/__init__.py new file mode 100644 index 0000000000..c7c2a85c06 --- /dev/null +++ b/src/diffusers/pipelines/score_sde_ve/__init__.py @@ -0,0 +1 @@ +from .pipeline_score_sde_ve import ScoreSdeVePipeline diff --git a/src/diffusers/pipelines/pipeline_score_sde_ve.py b/src/diffusers/pipelines/score_sde_ve/pipeline_score_sde_ve.py similarity index 100% rename from src/diffusers/pipelines/pipeline_score_sde_ve.py rename to src/diffusers/pipelines/score_sde_ve/pipeline_score_sde_ve.py diff --git a/src/diffusers/pipelines/score_sde_vp/__init__.py b/src/diffusers/pipelines/score_sde_vp/__init__.py new file mode 100644 index 0000000000..40660ec662 --- /dev/null +++ b/src/diffusers/pipelines/score_sde_vp/__init__.py @@ -0,0 +1 @@ +from .pipeline_score_sde_vp import ScoreSdeVpPipeline diff --git a/src/diffusers/pipelines/pipeline_score_sde_vp.py b/src/diffusers/pipelines/score_sde_vp/pipeline_score_sde_vp.py similarity index 100% rename from src/diffusers/pipelines/pipeline_score_sde_vp.py rename to src/diffusers/pipelines/score_sde_vp/pipeline_score_sde_vp.py diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py index 3ce544fab8..d659d7ca72 100755 --- a/tests/test_modeling_utils.py +++ b/tests/test_modeling_utils.py @@ -47,7 +47,7 @@ from diffusers import ( ) from diffusers.configuration_utils import ConfigMixin from diffusers.pipeline_utils import DiffusionPipeline -from diffusers.pipelines.pipeline_bddm import DiffWave +from diffusers.pipelines.bddm.pipeline_bddm import DiffWave from diffusers.testing_utils import floats_tensor, slow, torch_device @@ -1005,11 +1005,13 @@ class PipelineTesterMixin(unittest.TestCase): bddm = BDDMPipeline(model, noise_scheduler) # check if the library name for the diffwave moduel is set to pipeline module - self.assertTrue(bddm.config["diffwave"][0] == "pipeline_bddm") + self.assertTrue(bddm.config["diffwave"][0] == "bddm") # check if we can save and load the pipeline with tempfile.TemporaryDirectory() as tmpdirname: bddm.save_pretrained(tmpdirname) _ = BDDMPipeline.from_pretrained(tmpdirname) # check if the same works using the DifusionPipeline class - _ = DiffusionPipeline.from_pretrained(tmpdirname) + bddm = DiffusionPipeline.from_pretrained(tmpdirname) + + self.assertTrue(bddm.config["diffwave"][0] == "bddm")