wip scheduler

huggingface · Dec 4, 2023 · 111be88 · 111be88
1 parent 0f55c17
commit 111be88
Show file tree

Hide file tree

Showing 6 changed files with 448 additions and 1 deletion.
diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
@@ -383,6 +383,7 @@
         [
             "FlaxDDIMScheduler",
             "FlaxDDPMScheduler",
+            "FlaxLCMScheduler",
             "FlaxDPMSolverMultistepScheduler",
             "FlaxEulerDiscreteScheduler",
             "FlaxKarrasVeScheduler",
@@ -705,6 +706,7 @@
         from .schedulers import (
             FlaxDDIMScheduler,
             FlaxDDPMScheduler,
+            FlaxLCMScheduler,
             FlaxDPMSolverMultistepScheduler,
             FlaxEulerDiscreteScheduler,
             FlaxKarrasVeScheduler,

diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py
@@ -86,6 +86,7 @@
     "diffusers": {
         "ModelMixin": ["save_pretrained", "from_pretrained"],
         "SchedulerMixin": ["save_pretrained", "from_pretrained"],
+        "FlaxSchedulerMixin": ["save_pretrained", "from_pretrained"],
         "DiffusionPipeline": ["save_pretrained", "from_pretrained"],
         "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"],
     },

diff --git a/src/diffusers/pipelines/stable_diffusion_xl/__init__.py b/src/diffusers/pipelines/stable_diffusion_xl/__init__.py
@@ -29,6 +29,7 @@
     _import_structure["pipeline_stable_diffusion_xl_img2img"] = ["StableDiffusionXLImg2ImgPipeline"]
     _import_structure["pipeline_stable_diffusion_xl_inpaint"] = ["StableDiffusionXLInpaintPipeline"]
     _import_structure["pipeline_stable_diffusion_xl_instruct_pix2pix"] = ["StableDiffusionXLInstructPix2PixPipeline"]
+    _import_structure["pipeline_flax_stable_diffusion_xl"] = ["FlaxStableDiffusionXLPipeline"]
 
 if is_transformers_available() and is_flax_available():
     from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState

diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py
@@ -28,6 +28,7 @@
     FlaxDPMSolverMultistepScheduler,
     FlaxLMSDiscreteScheduler,
     FlaxPNDMScheduler,
+    FlaxLCMScheduler
 )
 from ..pipeline_flax_utils import FlaxDiffusionPipeline
 from .pipeline_output import FlaxStableDiffusionXLPipelineOutput
@@ -49,7 +50,7 @@ def __init__(
         tokenizer_2: CLIPTokenizer,
         unet: FlaxUNet2DConditionModel,
         scheduler: Union[
-            FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler
+            FlaxDDIMScheduler, FlaxPNDMScheduler, FlaxLMSDiscreteScheduler, FlaxDPMSolverMultistepScheduler, FlaxLCMScheduler
         ],
         dtype: jnp.dtype = jnp.float32,
     ):

diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py
@@ -77,6 +77,7 @@
 else:
     _import_structure["scheduling_ddim_flax"] = ["FlaxDDIMScheduler"]
     _import_structure["scheduling_ddpm_flax"] = ["FlaxDDPMScheduler"]
+    _import_structure["scheduling_lcm_flax"] = ["FlaxLCMScheduler"]
     _import_structure["scheduling_dpmsolver_multistep_flax"] = ["FlaxDPMSolverMultistepScheduler"]
     _import_structure["scheduling_euler_discrete_flax"] = ["FlaxEulerDiscreteScheduler"]
     _import_structure["scheduling_karras_ve_flax"] = ["FlaxKarrasVeScheduler"]
@@ -165,6 +166,7 @@
         from .scheduling_ddim_flax import FlaxDDIMScheduler
         from .scheduling_ddpm_flax import FlaxDDPMScheduler
         from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
+        from .scheduling_lcm_flax import FlaxLCMScheduler
         from .scheduling_euler_discrete_flax import FlaxEulerDiscreteScheduler
         from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
         from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler