Refactoring Diffusers pipelines

optimum/exporters/neuron/__main__.py

@@ -598,12 +598,6 @@ def main_export(
 
     # Validate compiled model
     if do_validation is True:
-        if library_name == "diffusers":
-            # Do not validate vae encoder due to the sampling randomness
-            neuron_outputs.pop("vae_encoder")
-            models_and_neuron_configs.pop("vae_encoder", None)
-            output_model_names.pop("vae_encoder", None)
-
         try:
             validate_models_outputs(
                 models_and_neuron_configs=models_and_neuron_configs,

optimum/exporters/neuron/model_configs.py

@@ -724,7 +724,7 @@ def inputs(self) -> List[str]:
 
     @property
     def outputs(self) -> List[str]:
-        return ["latent_sample"]
+        return ["latent_parameters"]
 
     def generate_dummy_inputs(self, return_tuple: bool = False, **kwargs):
         dummy_inputs = super().generate_dummy_inputs(**kwargs)

optimum/exporters/neuron/utils.py

@@ -421,7 +421,7 @@ def get_submodels_for_export_stable_diffusion(
 
     # VAE Encoder
     vae_encoder = copy.deepcopy(pipeline.vae)
-    vae_encoder.forward = lambda sample: {"latent_sample": vae_encoder.encode(x=sample)["latent_dist"].sample()}
+    vae_encoder.forward = lambda sample: {"latent_parameters": vae_encoder.encode(x=sample)["latent_dist"].parameters}
     models_for_export.append((DIFFUSION_MODEL_VAE_ENCODER_NAME, vae_encoder))
 
     # VAE Decoder

optimum/neuron/pipelines/__init__.py

@@ -21,35 +21,19 @@
 _import_structure = {
     "transformers": ["pipeline"],
     "diffusers": [
-        "NeuronStableDiffusionPipelineMixin",
-        "NeuronStableDiffusionImg2ImgPipelineMixin",
-        "NeuronStableDiffusionInpaintPipelineMixin",
-        "NeuronStableDiffusionInstructPix2PixPipelineMixin",
-        "NeuronLatentConsistencyPipelineMixin",
-        "NeuronStableDiffusionControlNetPipelineMixin",
         "NeuronStableDiffusionXLPipelineMixin",
-        "NeuronStableDiffusionXLImg2ImgPipelineMixin",
-        "NeuronStableDiffusionXLInpaintPipelineMixin",
+        "NeuronStableDiffusionControlNetPipelineMixin",
         "NeuronStableDiffusionXLControlNetPipelineMixin",
     ],
 }
 
 if TYPE_CHECKING:
     from .diffusers import (
-        NeuronLatentConsistencyPipelineMixin,
         NeuronStableDiffusionControlNetPipelineMixin,
-        NeuronStableDiffusionImg2ImgPipelineMixin,
-        NeuronStableDiffusionInpaintPipelineMixin,
-        NeuronStableDiffusionInstructPix2PixPipelineMixin,
-        NeuronStableDiffusionPipelineMixin,
         NeuronStableDiffusionXLControlNetPipelineMixin,
-        NeuronStableDiffusionXLImg2ImgPipelineMixin,
-        NeuronStableDiffusionXLInpaintPipelineMixin,
         NeuronStableDiffusionXLPipelineMixin,
     )
-    from .transformers import (
-        pipeline,
-    )
+    from .transformers import pipeline
 else:
     import sys
 

optimum/neuron/pipelines/diffusers/__init__.py

@@ -12,14 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 from .pipeline_controlnet import NeuronStableDiffusionControlNetPipelineMixin
 from .pipeline_controlnet_sd_xl import NeuronStableDiffusionXLControlNetPipelineMixin
-from .pipeline_latent_consistency_text2img import NeuronLatentConsistencyPipelineMixin
-from .pipeline_stable_diffusion import NeuronStableDiffusionPipelineMixin
-from .pipeline_stable_diffusion_img2img import NeuronStableDiffusionImg2ImgPipelineMixin
-from .pipeline_stable_diffusion_inpaint import NeuronStableDiffusionInpaintPipelineMixin
-from .pipeline_stable_diffusion_instruct_pix2pix import NeuronStableDiffusionInstructPix2PixPipelineMixin
-from .pipeline_stable_diffusion_xl import NeuronStableDiffusionXLPipelineMixin
-from .pipeline_stable_diffusion_xl_img2img import NeuronStableDiffusionXLImg2ImgPipelineMixin
-from .pipeline_stable_diffusion_xl_inpaint import NeuronStableDiffusionXLInpaintPipelineMixin
+from .pipeline_utils import NeuronStableDiffusionXLPipelineMixin